freerafiki / multi_part_assembly

Multi-Part Assembly Benchmark Implemented in PyTorch

Multi Part Assembly

Introduction

This is an open-source 3D shape assembly codebase based on PyTorch + PyTorch-Lightning.

Major Features

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.

Installation

Please refer to install.md for step-by-step guidance on how to install the packages and prepare the data.

Config System

To learn about the config system used in this codebase, please refer to the config.md.

Quick Start

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.

Tutorial

We explain our code design in tutorial.md. Please read it before modifying the codebase or implementing your new algorithms.

Benchmark

We benchmark the baselines and report them in model.md. We also present detailed instructions on how to reproduce our results.

License

This project is released under the MIT license.

Acknowledgement

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.