Home

Awesome

Logo

MultiBodySync: Multi-Body Segmentation and Motion Estimation via 3D Scan Synchronization

This repository contains a PyTorch implementation of the above paper. It will be presented at CVPR2021 as an oral.

Introduction

Logo

MultiBodySync is an end-to-end trainable multi-body motion segmentation and rigid registration framework for multiple input 3D point clouds with the following features:

Usage

Please install Pytorch (>=1.6) and run the following code (e.g. in a notebook), then you get our model ready for inference:

import torch
# Load model
my_model = torch.hub.load('huangjh-pub/multibody-sync:public', 'model_articulated', pretrained=True)
my_model.cuda().eval()
# Perform inference, data has to be (1, K, N, 3) cuda tensor.
with torch.no_grad():
    my_model.forward(data)

You may also need an example input (e.g. the one under assets/laptop.npy) to feed into the model. We suggest you normalize your point cloud (preferably with -Y axis up) beforehand to get the best result.

Dependencies

We suggest to use Anaconda to manage your environment. Following is the suggested way to install the dependencies:

# Create a new conda environment
conda create -n mbs python=3.8
conda activate mbs

# Install pytorch
conda install pytorch==1.6.0 cudatoolkit=10.2 -c pytorch

# Install other packages
pip install -r requirements.txt

For domestic users please consider using a mirror if there are connection problems.

Dataset

Each dataset is organized in the following structure:

<dataset-name>/
    ├ meta.json
    └ data/
        ├ 000000.npz
        ├ 000001.npz
        └ ...

After downloading the dataset, please set the paths in the corresponding yaml config files to the root of the dataset folder, i.e., <dataset-name>/.

Articulated Objects

Solid Objects

DynLab Dataset

The raw dataset can be downloaded Here. The files are organized as 8 scenes x 8 configurations = 64 (point cloud, pose) tuples, each of which is formatted as:

Training and Test

Please use the following commands for training. We suggest to train the flow network and mot network simultaneously and then train conf network after flow is fully converged.

# Train flow network
python train.py config/articulated-flow.yaml
# Train mot network
python train.py config/articulated-mot.yaml
# Train conf network
python train.py config/articulated-conf.yaml

Then the entire pipeline can be tuned end-to-end using the following:

python train.py config/articulated-full.yaml

After training, run the following to test your trained model:

python test.py config/articulated-full.yaml

Pre-trained models

Please download the corresponding trained weights for articulated objects or solid objects and extract the weights to ./ckpt/articulated-full/best.pth.tar.

For solid objects, simply do %s/articulated/solid/g.

Citation

The paper will not appear in the proceedings before the conference actually takes place. For now you may use the following bibtex:

@article{huang2021multibodysync,
  title={MultiBodySync: Multi-Body Segmentation and Motion Estimation via 3D Scan Synchronization},
  author={Huang, Jiahui and Wang, He and Birdal, Tolga and Sung, Minhyuk and Arrigoni, Federica and Hu, Shi-Min and Guibas, Leonidas},
  journal={arXiv preprint arXiv:2101.06605},
  year={2021}
}

Reference code