This work proposes exploiting equivariant learning from 3D point clouds to improve registration robustness. We propose SE3ET, an SE(3)-equivariant registration framework that employs equivariant point convolution and equivariant transformer design to learn expressive and robust geometric features.
@article{lin2024se3et,
title={SE3ET: SE(3)-Equivariant Transformer for Low-Overlap Point Cloud Registration},
author={Lin, Chien Erh and Zhu, Minghan and Ghaffari, Maani},
journal={IEEE Robotics and Automation Letters},
year={2024},
publisher={IEEE}
}
Please use the following command for installation.
Build docker environment or conda environment
- docker
cd docker
docker build --tag umcurly/geotransformer .
bash build_docker_container.bash [container_name]
- conda
conda create -n se3et python==3.8
conda activate se3et
pip install torch==1.7.1+cu110 -f https://download.pytorch.org/whl/torch_stable.html
pip install -r requirements.txt
Build extension package (might need root access)
python setup.py build develop
cd geotransformer/modules/e2pn/vgtk
python setup.py build_ext -i
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.
Please download the weights on Google Drive and put them in weights directory.
TL;DR: We suggest you to use SE3ET-E if the input point cloud is small (<10k points) and use SE3ET-I if the input point cloud is large (>10k points).
In the paper, we propose two transformer configurations. SE3ET-E contains all the equivariant and invariant self-attention and cross-attention we proposed, and SE3ET-I contains equivariant self-attention and invariant cross-attention. The first one provides better performance on low-overlap and random rotation cases but higher on memory consumption. The latter one maintain good performance while keeping lower memory consumption. We suggest you use SE3ET-E if the input point cloud is small (<10k points) and use SE3ET-I if the input point cloud is large (>10k points). We also provide SE3ET-E2 and SE3ET-I2 which the feature size is 2x smaller to boost computation performance. Please check the paper for their performance comparison.
If you want to test on your data, direct to either one of the experiment folder, and then run the demo code.
cd experiments/se3eti.3dmatch
python demo.py --src_file=<point_cloud_1_path> --ref_file=<point_cloud_2_path> --gt_file=<ground_truth_transformation_file> --weights=../../weights/se3eti.3dmatch.pth.tar
The dataset can be downloaded from PREDATOR. Check out PREDATOR and GeoTransformer repositories for preparing data.
The code for 3DMatch is in experiments/se3ete.3dmatch for SE3ET-E and experiments/se3eti.3dmatch for SE3ET-I. Inside the folder, change config.py accordingly. Then, use the following command for training.
CUDA_VISIBLE_DEVICES=0 python trainval.py
Command to generate features from point clouds
CUDA_VISIBLE_DEVICES=0 python test.py --snapshot=../../weights/se3ete.3dmatch.pth.tar --benchmark=3DMatch
CUDA_VISIBLE_DEVICES=0 python test.py --snapshot=../../weights/se3ete.3dmatch.pth.tar --benchmark=3DLoMatch
Command to run feature matching and registration
CUDA_VISIBLE_DEVICES=0 python eval.py --num_corr=5000 --benchmark=3DMatch --method=ransac
or
CUDA_VISIBLE_DEVICES=0 python eval.py --benchmark=3DMatch --method=lgr
Direct to folders experiments/se3ete.3dmatch.evalrot for SE3ET-E or experiments/se3eti.3dmatch.evalrot for SE3ET-I and follow the testing commands using the same pretrained weight. In these folders, we changed the following code to test the algorithm with random rotation of the input.
- In
dataset.py,rotated=True,is added in the configuration oftest_dataset = ThreeDMatchPairDataset()function. - In
eval.py,transform = gt_logs[gt_index]['transform']is commentted out. We are loading the transformation from data_dict constructed in dataloader.
Download the data from the Kitti official website into data/Kitti and run data/Kitti/downsample_pcd.py to generate the data.
The code for Kitti is in experiments/se3eti.kitti. Use the following command for training.
CUDA_VISIBLE_DEVICES=0 python trainval.pyUse the following command for testing.
CUDA_VISIBLE_DEVICES=0 python test.py --snapshot=../../weights/se3eti.kitti.pth.tar
CUDA_VISIBLE_DEVICES=0 python eval.py --method=lgrDirect to folder experiments/se3eti.kitti.evalrot and follow the testing commands using the same pretrained weight. In dataset.py, the following configuration is added to the test_dataset = OdometryKittiPairDataset() function to test the algorithm with random rotation of the input.
use_augmentation=cfg.train.use_augmentation,
augmentation_noise=cfg.train.augmentation_noise,
augmentation_min_scale=cfg.train.augmentation_min_scale,
augmentation_max_scale=cfg.train.augmentation_max_scale,
augmentation_shift=cfg.train.augmentation_shift,
augmentation_rotation=cfg.train.augmentation_rotation,
In the generalization test, we evaluate the SE3ET-I2 model that trained on 3DMatch on (scaled) KITTI point cloud pairs.
Direct to experiment/se3eti2.3dmatch.evalkitti, and run
CUDA_VISIBLE_DEVICES=0 python test.py --snapshot=../../weights/se3eti.kitti.pth.tar
CUDA_VISIBLE_DEVICES=0 python eval.py --method=lgrWe thank the authors for open sourcing their methods, the code is heavily borrowed from GeoTransformer and E2PN.