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]
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.
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/
.
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.
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.