Machine learning improves pose estimation of LEGO parts in images

BrickFlow: Connectivity-Guided Brick Reconstruction

Computer Vision and Pattern RecognitionGraphics

Summary

Figuring out how LEGO pieces fit together just by looking at a picture is very hard because the pieces have to connect physically in the right way. The authors created a machine learning model that learns from many fake LEGO builds how parts usually fit together. They then fine-tuned their model to understand real images with very little extra data. To make sure pieces actually connect well, they added a new method that guides the model’s guesses towards connected assemblies. This approach helps get better and more accurate LEGO reconstructions from images.

What this means in practice

  • For robotics engineers: Enable robots to assemble LEGO structures by interpreting images and placing parts with physically correct connections.
  • For ar/vr developers: Create virtual LEGO assembly guides that understand real-world LEGO models from images to assist users in building.

Authors

Peter Kulits, Cordelia Schmid

Abstract

Given a set of LEGO parts and an image of an assembled object, our task is to infer the pose of each part to assemble the object. In contrast with traditional 3D reconstruction, doing so requires satisfying discrete physical constraints, which provide structure but render the problem combinatorial and highly sensitive to small pose errors. To learn a prior over how parts fit together, we train a set-conditioned SE(3) flow model on procedurally generated assemblies. We then finetune this model to inject image conditioning, using 1% as much data. However, while the learned flow captures global structure, it does not enforce connectivity. To address this, we introduce an analytic connectivity-field guidance to pull predictions toward connected assemblies at test time. Through this, we demonstrate improved build quality with model size, pretraining, and guidance. We make our models and data available for research purposes at https://kulits.github.io/BrickFlow/