This repository represents the official implementation of the ECCV'2022 paper: "Improving RGB-D Point Cloud Registration by Learning Multi-scale Local Linear Transformation" by Ziming Wang*, Xiaoliang Huo*, Zhenghao Chen, Jing Zhang, Lu Sheng† and Dong Xu (*equal contributions, †corresponding author)

If you have any questions, please let us know:

• Ziming Wang {[email protected]}
• Xiaoliang Huo {[email protected]}

In this work, we propose a new Geometry-Aware Visual Feature Extractor (GAVE) that employs multi-scale local linear transformation to progressively fuse these two modalities, where the geometric features from the depth data act as the geometry-dependent convolution kernels to transform the visual features from the RGB data. The resultant visual-geometric features are in canonical feature spaces with alleviated visual dissimilarity caused by geometric changes, by which more reliable correspondence can be achieved. The proposed GAVE module can be readily plugged into recent RGB-D point cloud registration framework.

We train our module under two different setups:

• Training on the 3D Match dataset for 14 epochs, and testing on the ScanNet test set.
• Training on the ScanNet dataset for 1 epoch, and testing on the ScanNet test set.

The overall results are shown in the chart below along with the checkpoints' download links:

Here are several visualization examples of our method comparing to our baseline and Ground Truth:

This code has been trained/tested on:

• Python 3.6.13, PyTorch 1.7.1, CUDA 11.0.3, gcc 9.3.0, Tesla V100-PCIE-32GB

# create a conda environment and activate it
conda create --name GAVE python=3.10
conda activate GAVE

# install pytorch (any version that match your CUDA version)
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
conda install matplotlib tensorboard

# install pytorch3d
conda install pytorch3d -c pytorch3d

# install open3d
pip -m pip install open3d

# install Minkowski Engine 
git clone https://github.com/NVIDIA/MinkowskiEngine
cd MinkowskiEngine
python setup.py install --blas=openblas --blas_include_dirs=${CONDA_PREFIX}/include

# install other dependencies
python -m pip install nibabel opencv-python easydict pre-commit

If any error happens when installing the Minkowski Engine, please follow the official instruction.

Following the UnsupervisedR&R, we ues two datasets for training in our work: 3DMatch and ScanNet, and evaluate only in ScanNet test set.

For the download and pre-processing procedure, please refer to UR&R's instruction.

After downloading the datasets, make sure to update the paths in GAVE/datasets/builder.py with the dataaset root directories.

You can modify the settings of models in GAVE/configs/config.py and appoint the GPU in train.py:

# Training on 3DMatch dataset
python -u train.py --mode train --config_path 3DMatch.yaml
# Training on ScanNet dataset
python -u train.py --mode train --config_path ScanNet.yaml

You can evaulate checkpoints using the following command:

python -u evaluate.py mine --checkpoint ckpt_path/ckpt_name.pkl --progress_bar --boost_alignment

This repo heavily benefits from the UnsupervisedR&R. We would like to thank Mohamed El for his excellent work.

@inproceedings{wang2022improving,
  title={Improving rgb-d point cloud registration by learning multi-scale local linear transformation},
  author={Wang, Ziming and Huo, Xiaoliang and Chen, Zhenghao and Zhang, Jing and Sheng, Lu and Xu, Dong},
  booktitle={Computer Vision--ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23--27, 2022, Proceedings, Part XXXII},
  pages={175--191},
  year={2022},
  organization={Springer}
}