Translating a Visual LEGO Manual to a Machine-Executable Plan

This is the PyTorch implementation for the paper:

Translating a Visual LEGO Manual to a Machine-Executable Plan

Ruocheng Wang, Yunzhi Zhang, Jiayuan Mao, Chin-Yi Cheng, Jiajun Wu

In European Conference on Computer Vision (ECCV) 2022

[project]

Installation

Run the following commands to install necessary dependencies.

conda create -n lego_release python=3.9.12
conda activate lego_release
pip -r requirements.txt

You may need to manually install pytorch3d 0.5.0 according to this doc.

Evaluation

Download the evaluation datasets and model checkpoints from here, and unzip them under the root directory of the code. Then simply run

bash scripts/eval/eval_all.sh

from the root directory. Results will be saved to results/.

Training

To train our model from scratch, first download the training and validation datasets from here, and unzip them to data/datasets/synthetic_train and data/datasets/synthetic_val respectively.

After downloading the datasets, preprocess them by running

bash scripts/process_dataset.sh

Then run the script to train our model

bash scripts/train/train_mepnet.sh

You can add --wandb option in the training script for logging and visualization in wandb. We train our model on 4 Titan RTX GPUs for 5 days.

Acknowledgements

Some of our code is built on top of CycleGAN and CenterNet. If you encounter any problem, please don't hesitate to email me at rcwang@stanford.edu or open an issue.