CorrNet3D

Official implementation of our work as described in CorrNet3D: Unsupervised End-to-end Learning of Dense Correspondence for 3D Point Clouds (CVPR'21)

Note

2022-04: Since pytorch-lightning is updating, it's the most convenient for us to use the specific old version (download from their release) instead of the newest one. Version info: Python3.8, pytorch-lightning 1.1.6, pytorch 1.6.0, cudatoolkit 10.2, Cuda 10.2 (on RTX 2080 Ti) and Ubuntu 18.04.3 LTS (x86_64)

Prerequisite Installation

The code has been tested with Python3.8, pytorch-lightning 1.1.6 and Cuda 10.2:

conda create --name corrnet3d python=3.8
conda activate corrnet3d
pip install pytorch-lightning==1.1.6
pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl
pip install "git+git://github.com/erikwijmans/Pointnet2_PyTorch.git#egg=pointnet2_ops&subdirectory=pointnet2_ops_lib"
conda install torchvision torchaudio cudatoolkit=10.2 -c pytorch
pip install h5py
pip install tables
pip install matplotlib

Usage

Pre-trained Models

Download the pre-trained folder and put the folder under the corrnet3d folder.

Datasets (same .h5 for both non-rigid & rigid case)

We provide the .h5 dataset, inherited from the main train/test dataset from 3d-coded. You can download the dataset we used in the paper:

• train.h5
• test.h5

Additional dataset for SHREC (for non-rigid test only)

• tar file

Train & Test

To test on the whole testing set, run:

uncomment 'cli_main_test_()' in lit_corrnet3d_clean.py
python lit_corrnet3d_clean.py --gpus=3 --batch_size=1 --ckpt_user=lightning_logs/version_114/checkpoints/epoch=43-step=202399.ckpt --data_dir=./trainset.h5 --test_data_dir=./testset.h5

To train the network, run:

uncomment 'cli_main()' in lit_corrnet3d_clean.py
python lit_corrnet3d_clean.py --gpus=3 --batch_size=20 --data_dir=./trainset.h5 --test_data_dir=./testset.h5

Citation

Please cite this paper with the following bibtex:

@inproceedings{zeng2020corrnet3d,
    author={Zeng, Yiming and Qian, Yue and Zhu, Zhiyu and Hou, Junhui and Yuan, Hui and He, Ying},
    title={CorrNet3D: Unsupervised End-to-end Learning of Dense Correspondence for 3D Point Clouds},
    booktitle = {{IEEE/CVF} Conference on Computer Vision and Pattern Recognition (CVPR)},
    year      = {2021}
}