Equivariant Point Network (EPN)

This repository contains the code (in PyTorch) for Equivariant Point Network for 3D Point Cloud Analysis (CVPR'2021) by Haiwei Chen, Shichen Liu, Weikai Chen and Hao Li.

Contents

1. Introduction
2. Usage
3. Experiments
4. Contact

Introduction

EPN is a SE(3)-equivariant network model that is designed for deep point cloud analysis. The core of the architecture is the SE(3) Separable Convolution that combines two sequential, equivariant convolution layers to approximate convolution in the SE(3) space. With the incorporation of an attention mechanism, the EPN network model can be used to extract both SE(3) equivariant features and selectively pooled invariant features for various feature learning tasks.

Usage

The code has been tested on Python3.7, PyTorch 1.7.1 and CUDA (10.1). The module and additional dependencies can be installed with

cd vgtk
python setup.py install

Experiments

Datasets

The rotated Modelnet40 point cloud dataset is generated from the Aligned Modelnet40 subset and can be downloaded using this link.

The original 3DMatch training and evaluation dataset can be found here. We followed this repo to preprocess rgb frames into fused fragments and extract matching keypoints for training. The preprocessed data ready for training can be downloaded here (146GB). We also prepared to preprocessed 3DMatch evaluation dataset here (40GB), where local patches around testing keypoints have been precomputed.

Pretrained Model

Pretrained model can be downloaded using this link

Training

The following lines can be used for the training of each experiment

# modelnet classification
CUDA_VISIBLE_DEVICES=0 python run_modelnet.py experiment -d PATH_TO_MODELNET40
# modelnet shape alignment
CUDA_VISIBLE_DEVICES=0 python run_modelnet_rotation.py experiment -d PATH_TO_MODELNET40
# 3DMatch shape registration
CUDA_VISIBLE_DEVICES=0 python run_3dmatch.py experiment -d PATH_TO_3DMATCH

Evaluation

The following lines can be used for the evaluation of each experiment

# modelnet classification
CUDA_VISIBLE_DEVICES=0 python run_modelnet.py experiment -d PATH_TO_MODELNET40 -r PATH_TO_CKPT --run-mode eval
# modelnet shape alignment
CUDA_VISIBLE_DEVICES=0 python run_modelnet_rotation.py experiment -d PATH_TO_MODELNET40 -r PATH_TO_CKPT --run-mode eval
# 3DMatch shape registration
CUDA_VISIBLE_DEVICES=0 python run_3dmatch.py experiment -d PATH_TO_3DMATCH -r PATH_TO_CKPT --run-mode eval

Contact

Haiwei Chen: chw9308@hotmail.com Any discussions or concerns are welcomed!

Citation If you find our project useful in your research, please consider citing:

@article{chen2021equivariant,
  title={Equivariant Point Network for 3D Point Cloud Analysis},
  author={Chen, Haiwei and Liu, Shichen and Chen, Weikai and Li, Hao and Hill, Randall},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={14514--14523},
  year={2021}
}