Jigsaw: Learning to Assemble Multiple Fractured Objects

This repository contains a minimal PyTorch implementation of the paper "Jigsaw: Learning to Assemble Multiple Fractured Objects".

As we move forward, we are committed to enriching this codebase to further support research for assembly. Our plan involves expanding the range of supported datasets, encoders, models, visualization, and more supporting scripts! Please let us know which part should we prioritize.

Installation

This repository has been developed and tested with Ubuntu 18.04 and CUDA 11.0. To set up the required environment, follow these steps:

Create a new Anaconda environment named assembly:
```
conda env create -f environment.yaml
conda activate assembly
```
You may also use environment_simple.yaml for better adaptability, but it is not tested on our machine.
Install custom CUDA ops for Chamfer distance:
```
cd utils/chamfer
pip install -e .
```

Datasets

We provide support for the Breaking Bad Dataset. For more information about data processing, please refer to the dataset website. Please make sure you use the updated inner-face-free version, as our tests are all based on that version.

After processing the data, ensure that you have a folder named data with the following structure:

data
├── breaking_bad
│   ├── everyday
│   │   ├── BeerBottle
│   │   │   ├── ...
│   │   ├── ...
│   ├── everyday.train.txt
│   ├── everyday.val.txt
│   └── ...
└── ...

Only the everyday subset is necessary for training. If you want to test the artifact and other subsets, the structure should follow the same pattern.

Run the Experiment

For training, run

python train_matching.py --cfg path/to/your/yaml

Replace path/to/your/yaml with the path to your configuration file, for example:

python train_matching.py --cfg experiments/jigsaw_250e_cosine.yaml

The results will be stored to the directory results/MODEL_NAME/.

For evaluation, run

python eval_matching.py --cfg path/to/your/eval_yaml

Default configuration files are stored in the experiments/ directory, and you are encouraged to try your own configurations. If you discover a better configuration, please let us know by raising an issue or a pull request, and we will update the benchmark accordingly!

Pretrained Weights

To use the pretrained weights, download the weight file here and add the following configuration to your configuration file:

WEIGHT_FILE: path/to/your/weight_file.ckpt

Tutorial

We provide a tutorial to help you better understand each component of this code base.

Acknowledgement

We would like to express our gratitude to the authors of the following repositories, from which we referenced code:

Multi-part-assembly: breaking bad benchmark code.
ThinkMatch: an open sourced repository for deep graph matching algorithms.
Pointcept: providing point transformer structure.

We would also like to thank the authors of all the packages we use. We welcome any valuable suggestions for improving our repository.

Citation

If you find this repository useful in your research, please cite

@inproceedings{lu2023jigsaw,
    title={Jigsaw: Learning to Assemble Multiple Fractured Objects},
    author={Jiaxin Lu and Yifan Sun and Qixing Huang},
    booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
    year={2023},
    url={https://openreview.net/forum?id=OwpaO4w6K7}
}

Jiaxin-Lu/Jigsaw