Created by Xumin Yu*, Yongming Rao*, Ziyi Wang, Zuyan Liu, Jiwen Lu, Jie Zhou
Adapted by Ben Wagner
[arXiv] [Video] [Dataset] [Models] [supp]
This repository contains PyTorch implementation for PoinTr: Diverse Point Cloud Completion with Geometry-Aware Transformers (ICCV 2021 Oral Presentation).
PoinTr is a transformer-based model for point cloud completion. By representing the point cloud as a set of unordered groups of points with position embeddings, we convert the point cloud to a sequence of point proxies and employ a transformer encoder-decoder architecture for generation. We also propose two more challenging benchmarks ShapeNet-55/34 with more diverse incomplete point clouds that can better reflect the real-world scenarios to promote future research.
The most successful model for ALPhA's purposes so far has been SnowflakeNet, which is included in this repository.
- PyTorch >= 1.7.0
- python >= 3.7
- CUDA >= 9.0
- GCC >= 4.9
- torchvision
- timm
- open3d
- tensorboardX
pip install -r requirements.txt NOTE: PyTorch >= 1.7 and GCC >= 4.9 are required. My environment has PyTorch = 2.0.1 and Python = 3.10.11
# Chamfer Distance bash install.sh The solution for a common bug in chamfer distance installation can be found in Issue #6
# PointNet++ pip install "git+https://github.com/erikwijmans/Pointnet2_PyTorch.git#egg=pointnet2_ops&subdirectory=pointnet2_ops_lib" # GPU kNN pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl Note: If you still get ModuleNotFoundError: No module named 'gridding' or something similar then run these steps
1. cd into extensions/Module (eg extensions/gridding) 2. run `python setup.py install` That will fix the ModuleNotFoundError.
The details of the data used so far for this project can be found in the ALPhA #data Slack Channel. Simulated 22Mg data has been the primary training dataset.
To inference sample(s) with pretrained model
python tools/inference.py \ ${POINTR_CONFIG_FILE} ${POINTR_CHECKPOINT_FILE} \ [--experimental] \ [--save_img_path <dir>] \ [--n_imgs <number>] For example, inference 10 samples from MidCutSnowflake.yaml and save the results under ATTPCPoinTr/imgs/
python tools/inference.py \ cfgs/Mg22_Ne20pp_models/MidCutSnowflake.yaml path/to/ckpt.pth \ --save_img_path ./imgs/ \ --n_imgs=10 \ To evaluate a pre-trained PoinTr model on the Three Dataset with single GPU, run:
bash ./scripts/test.sh <GPU_IDS> \ --ckpts <path> \ --config <config> \ --exp_name <name> Test the PoinTr pretrained model on the PCN benchmark:
bash ./scripts/test.sh 0 \ --ckpts ./pretrained/PoinTr_PCN.pth \ --config ./cfgs/PCN_models/PoinTr.yaml \ --exp_name example Test the PoinTr pretrained model on ShapeNet55 benchmark (easy mode):
bash ./scripts/test.sh 0 \ --ckpts ./pretrained/PoinTr_ShapeNet55.pth \ --config ./cfgs/ShapeNet55_models/PoinTr.yaml \ --mode easy \ --exp_name example Test the PoinTr pretrained model on the KITTI benchmark:
bash ./scripts/test.sh 0 \ --ckpts ./pretrained/PoinTr_KITTI.pth \ --config ./cfgs/KITTI_models/PoinTr.yaml \ --exp_name example CUDA_VISIBLE_DEVICES=0 python KITTI_metric.py \ --vis <visualization_path> To train a point cloud completion model from scratch, run:
# Use DataParallel (DP) bash ./scripts/train.sh <GPUIDS> \ --config <config> \ --exp_name <name> \ [--resume] \ [--start_ckpts <path>] \ [--val_freq <int>] Resume a checkpoint:
CUDA_VISIBLE_DEVICES=0,1 bash ./scripts/dist_train.sh 2 13232 \ --config ./cfgs/PCN_models/PoinTr.yaml \ --exp_name example --resume Finetune a PoinTr on PCNCars
CUDA_VISIBLE_DEVICES=0,1 bash ./scripts/dist_train.sh 2 13232 \ --config ./cfgs/KITTI_models/PoinTr.yaml \ --exp_name example \ --start_ckpts ./weight.pth Train a PoinTr model with a single GPU:
bash ./scripts/train.sh 0 \ --config ./cfgs/KITTI_models/PoinTr.yaml \ --exp_name example We also provide the Pytorch implementation of several baseline models including GRNet, PCN, TopNet and FoldingNet. For example, to train a GRNet model on ShapeNet-55, run:
CUDA_VISIBLE_DEVICES=0,1 bash ./scripts/dist_train.sh 2 13232 \ --config ./cfgs/ShapeNet55_models/GRNet.yaml \ --exp_name example MIT License
Our code is inspired by GRNet and mmdetection3d and SnowflakeNet.
If you find our work useful in your research, please consider citing:
@inproceedings{yu2021pointr, title={PoinTr: Diverse Point Cloud Completion with Geometry-Aware Transformers}, author={Yu, Xumin and Rao, Yongming and Wang, Ziyi and Liu, Zuyan and Lu, Jiwen and Zhou, Jie}, booktitle={ICCV}, year={2021} }