Awesome
This is the official implementation of the paper Leveraging SE(3)-Equivariance for Learning 3D Geometric Shape Assembly (ICCV 2023)
Environment Installation
Our environment dependencies are included in the "BreakingBad\multi_part_assembly.egg-info\requires.txt" file. You can easily install them when you create a new conda environment.
Dataset Generation
For Breaking Bad dataset, please download the dataset following https://breaking-bad-dataset.github.io/.
For Geometric Shape Mating dataset, please generate the dataset following https://github.com/pairlab/NSM. Besides, you can download the dataset we generate in https://www.dropbox.com/scl/fi/xigal4s4xmksie4ihm7su/ShapeNet_0103.zip?rlkey=54f9w1fcnc2fjm5xmijpmkzf5&dl=0.
Train Models
To train models for Breaking Bad dataset, pleas run
cd BreakingBad
python scripts/train.py --cfg_file configs/vnn/vnn-everyday.py
To train models for Geometric Shape Mating dataset, pleas run
cd NSM
python script/train_eqv_CR.py --cfg_file train_vnn_pn.yml
Evaluate Models
The evaluation of models in Breaking Bad dataset is automatically performed during training.
To evaluate models for Geometric Shape Mating dataset, pleas run
cd NSM
python script/eval_eqv_CR.py --cfg_file eval_vnn_pn.yml
Citation
If you find this paper useful, please consider citing:
@InProceedings{Wu_2023_ICCV,
author = {Wu, Ruihai and Tie, Chenrui and Du, Yushi and Zhao, Yan and Dong, Hao},
title = {Leveraging SE(3) Equivariance for Learning 3D Geometric Shape Assembly},
booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
month = {October},
year = {2023},
pages = {14311-14320}
}
Contact
If you have any questions, please feel free to contact Ruihai Wu at wuruihai_at_pku_edu_cn and Chenrui Tie at crtie_at_pku_edu_cn