Group Equivariant Network for 3D Point Cloud Registration

PyTorch implementation of the paper "Group Equivariant Network for 3D Point Cloud Registration"

Installation

Please first install "torch" and "torch_geometric". Then you can use the following command for installation.

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.11.0, CUDA 11.3.

Data preparation

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
                          |                    |--...
                          |--...

Training

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

CUDA_VISIBLE_DEVICES=0 python trainval.py

Testing

Use the following command for testing.

CUDA_VISIBLE_DEVICES=0 python test.py --snapshot=../../output/3dmatch.stage4.gse.k3.max.oacl.stage2.sinkhorn/snapshots/snapshot.pth.tar --benchmark=3DMatch
CUDA_VISIBLE_DEVICES=0 python eval.py --benchmark=3DMatch --method=lgr

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