PyTorch implementation of the paper:

Geometric Transformer for Fast and Robust Point Cloud Registration.

Zheng Qin, Hao Yu, Changjian Wang, Yulan Guo, Yuxing Peng, and Kai Xu.

We study the problem of extracting accurate correspondences for point cloud registration. Recent keypoint-free methods bypass the detection of repeatable keypoints which is difficult in low-overlap scenarios, showing great potential in registration. They seek correspondences over downsampled superpoints, which are then propagated to dense points. Superpoints are matched based on whether their neighboring patches overlap. Such sparse and loose matching requires contextual features capturing the geometric structure of the point clouds. We propose Geometric Transformer to learn geometric feature for robust superpoint matching. It encodes pair-wise distances and triplet-wise angles, making it robust in low-overlap cases and invariant to rigid transformation. The simplistic design attains surprisingly high matching accuracy such that no RANSAC is required in the estimation of alignment transformation, leading to $100$ times acceleration. Our method improves the inlier ratio by $17% \sim 30%$ and the registration recall by over $7%$ on the challenging 3DLoMatch benchmark.

2022.03.30: Code and pretrained models on KITTI and ModelNet40 release.

2022.03.29: This work is selected for an ORAL presentation at CVPR 2022.

2022.03.02: This work is accepted by CVPR 2022. Code and Models on ModelNet40 and KITTI will be released soon.

2022.02.15: Paper is available at arXiv.

2022.02.14: Code and pretrained model on 3DMatch/3DLoMatch release.

Please use the following command for installation.

# It is recommended to create a new environment
conda create -n geotransformer python==3.8
conda activate geotransformer

# [Optional] If you are using CUDA 11.0 or newer, please install `torch==1.7.1+cu110`
pip install torch==1.7.1+cu110 -f https://download.pytorch.org/whl/torch_stable.html

# Install packages and other dependencies
pip install -r requirements.txt
python setup.py build develop

Code has been tested with Ubuntu 20.04, GCC 9.3.0, Python 3.8, PyTorch 1.7.1, CUDA 11.1 and cuDNN 8.1.0.

We provide pre-trained weights in the release page. Please download the latest weights and put them in weights directory.

The dataset can be downloaded from PREDATOR. The data should be organized as follows:

--data--3DMatch--metadata
              |--data--train--7-scenes-chess--cloud_bin_0.pth
                    |      |               |--...
                    |      |--...
                    |--test--7-scenes-redkitchen--cloud_bin_0.pth
                          |                    |--...
                          |--...

The code for 3DMatch is in experiments/geotransformer.3dmatch.stage4.gse.k3.max.oacl.stage2.sinkhorn. Use the following command for training.

CUDA_VISIBLE_DEVICES=0 python trainval.py

Use the following command for testing.

# 3DMatch
CUDA_VISIBLE_DEVICES=0 ./eval.sh EPOCH 3DMatch
# 3DLoMatch
CUDA_VISIBLE_DEVICES=0 ./eval.sh EPOCH 3DLoMatch

EPOCH is the epoch id.

We also provide pretrained weights in weights, use the following command to test the pretrained weights.

CUDA_VISIBLE_DEVICES=0 python test.py --snapshot=../../weights/geotransformer-3dmatch.pth.tar --benchmark=3DMatch
CUDA_VISIBLE_DEVICES=0 python eval.py --benchmark=3DMatch --method=lgr

Replace 3DMatch with 3DLoMatch to evaluate on 3DLoMatch.

Download the data from the Kitti official website into data/Kitti and run data/Kitti/downsample_pcd.py to generate the data. The data should be organized as follows:

--data--Kitti--metadata
            |--sequences--00--velodyne--000000.bin
            |              |         |--...
            |              |...
            |--downsampled--00--000000.npy
                         |   |--...
                         |--...

The code for Kitti is in experiments/geotransformer.kitti.stage5.gse.k3.max.oacl.stage2.sinkhorn. Use the following command for training.

CUDA_VISIBLE_DEVICES=0 python trainval.py

Use the following command for testing.

CUDA_VISIBLE_DEVICES=0 ./eval.sh EPOCH

EPOCH is the epoch id.

We also provide pretrained weights in weights, use the following command to test the pretrained weights.

CUDA_VISIBLE_DEVICES=0 python test.py --snapshot=../../weights/geotransformer-kitti.pth.tar
CUDA_VISIBLE_DEVICES=0 python eval.py --method=lgr

Download the data and run data/ModelNet/split_data.py to generate the data. The data should be organized as follows:

--data--ModelNet--modelnet_ply_hdf5_2048--...
               |--train.pkl
               |--val.pkl
               |--test.pkl

The code for ModelNet is in experiments/geotransformer.modelnet.rpmnet.stage4.gse.k3.max.oacl.stage2.sinkhorn. Use the following command for training.

CUDA_VISIBLE_DEVICES=0 python trainval.py

Use the following command for testing.

CUDA_VISIBLE_DEVICES=0 python test.py --test_iter=ITER

ITER is the iteration id.

We also provide pretrained weights in weights, use the following command to test the pretrained weights.

CUDA_VISIBLE_DEVICES=0 python test.py --snapshot=../../weights/geotransformer-modelnet.pth.tar

As the point clouds usually have different sizes, we organize them in the pack mode. This causes difficulty for batch training as we need to convert the data between batch mode and pack mode frequently. For this reason, we limit the batch size to 1 per GPU at this time and support batch training via DistributedDataParallel. Use torch.distributed.launch for multi-gpu training:

CUDA_VISIBLE_DEVICES=GPUS python -m torch.distributed.launch --nproc_per_node=NGPUS trainval.py

Note that the learning rate is multiplied by the number of GPUs by default as the batch size increased. In our experiments, multi-gpu training slightly improves the performance.

We evaluate GeoTransformer on the standard 3DMatch/3DLoMatch benchmarks as in PREDATOR.

We evaluate GeoTransformer on the standard Kitti benchmark as in PREDATOR.

We evaluate GeoTransformer on ModelNet with two settings:

1. Standard setting: [0, 45] rotation, [-0.5, 0.5] translation, gaussian noise clipped to 0.05.
2. Full-range setting: [0, 180] rotation, [-0.5, 0.5] translation, gaussian noise clipped to 0.05.

We remove symmetric classes and use the data augmentation in RPMNet which is more difficult than PRNet.

To test on your own data, the recommended way is to implement a Dataset as in geotransformer.dataset.registration.threedmatch.dataset.py. Each item in the dataset is a dict contains at least 5 keys: ref_points, src_points, ref_feats, src_feats and transform.

We also provide a demo script to quickly test our pre-trained model on your own data in experiments/geotransformer.3dmatch.stage4.gse.k3.max.oacl.stage2.sinkhorn/demo.py. Use the following command to run the demo:

CUDA_VISIBLE_DEVICES=0 python demo.py --src_file=../../data/demo/src.npy --ref_file=../../data/demo/ref.npy --gt_file=../../data/demo/gt.npy --weights=../../weights/geotransformer-3dmatch.pth.tar

Change the arguments src_file, ref_file and gt to your own data, where src_file and ref_file are numpy files containing a np.ndarray in shape of Nx3, and gt_file is a numpy file containing a 4x4 transformation matrix. Note that you should scale your data to match the voxel size in 3DMatch (2.5cm).

@inproceedings{qin2022geometric,
    title={Geometric Transformer for Fast and Robust Point Cloud Registration},
    author={Zheng Qin and Hao Yu and Changjian Wang and Yulan Guo and Yuxing Peng and Kai Xu},
    booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month={June},
    year={2022},
    pages={11143-11152}
}

• D3Feat
• PREDATOR
• RPMNet
• CoFiNet
• huggingface-transformer
• SuperGlue