This is an open-source 3D shape assembly codebase based on PyTorch + PyTorch-Lightning.
We mainly focus on vision based 3D shape assembly task, which takes in point clouds of multiple parts from an object, and predicts their rotations and translations of the correct assembly.
- Support datasets:
- PartNet (semantic assembly)
- Breaking-Bad (geometric assembly)
- Support models:
- Global, LSTM, DGL (NeurIPS 2020)
- RGL-NET (WACV 2022)
- Transformer-based (designed by us)
We carefully benchmark the models to match their performance in their original papers, which you can easily extend to new datasets. You can also leverage our codebase to develop your new shape assembly algorithms.
Please refer to install.md for step-by-step guidance on how to install the packages and prepare the data.
To learn about the config system used in this codebase, please refer to the config.md.
We provide detailed usage of the codebase in usage.md. You can train, test and visualize the results using our provided config files, or develop your new methods.
We explain our code design in tutorial.md. Please read it before modifying the codebase or implementing your new algorithms.
This project is released under the MIT license.
We thank the authors of the following repos for open-sourcing their wonderful works:
- MMCV, MMDetection3D: general project structure
- PyTorch3D: 3D transformation implementation
- DGL: shape assembly model code
- RGL-NET: shape assembly model code
We also thank the authors of all the packages we use. We appreciate all the contributors as well as users who give valuable feedbacks. We wish this codebase could serve as a benchmark and a flexible toolkit for researchers to re-implement existing algorithms and develop their own new shape assembly methods.