This is the official repo for PyTorch implementation of paper "SPoVT: Semantic-Prototype Variational Transformer for Dense Point Cloud Semantic Completion", NeurIPS 2022.
Paper | Project page | Video
- OS: Linux Ubuntu 20.04
- GPU: NVIDIA RTX A6000 or NVIDIA RTX 3090
- Python: 3.8.13
- CUDA Toolkit: 11.6
- GCC/G++: 7.2.0
- Python package manager
conda
Download datasets from here and put datasets under folder data/
.
Install all python packages for training and evaluation with conda environment by following commands:
conda create -n spovt python=3.8
conda activate spovt
conda install pytorch=1.13.0 torchvision pytorch-cuda=11.6 -c pytorch -c nvidia
conda install -c fvcore -c iopath -c conda-forge fvcore iopath
pip install "git+https://github.com/facebookresearch/pytorch3d.git"
pip install -r requirements.txt
Download pretrained model weights from here for evaluation without training from scratch. Put checkpoint files under folder ckpts/
.
We provide training commands for only car
experiment in the script file train_car.sh
, and evaluation commands for all experiment in the script file test.sh
.
Please note that the pretrained checkpoints have undergone extensive training, so it may require more epochs to achieve the performance reported in the paper.
Train the model on car
category by running:
bash train_car.sh
The model is trained progressively in 4 stages.
# Reconstruction Training (autoencoder)
python ./main.py -record records/car_pretrain.txt -info "car pretrain" -bs 64 -lr 1e-3 -epoch 200 -interval 500 -cuda 0 -save ckpts/car_pretrain.pth -mode pretrain -cat car
# Coarse Completion Training
python ./main.py -record records/car_train.txt -info "car train" -bs 64 -lr 1e-4 -epoch 200 -interval 500 -cuda 0 -save ckpts/car_train.pth -mode train -cat car -load ckpts/car_pretrain.pth
# Fine Completion Training (Stage 1)
python ./main.py -record records/car_refine.txt -info "car refine" -bs 64 -lr 1e-4 -epoch 200 -interval 500 -cuda 0 -save ckpts/car_refine.pth -mode refine -cat car -load ckpts/car_train.pth
# Fine Completion Training (Stage 2)
python ./main.py -record records/car_refine2.txt -info "car refine2" -bs 64 -lr 1e-4 -epoch 200 -interval 500 -cuda 0 -save ckpts/car_refine2.pth -mode refine2 -cat car -load ckpts/car_refine.pth
Generate complete point cloud and render segmented images by running:
bash test.sh
The rendered images are saved under folder seg_image
, and the generated point cloud are saved under folder visualize
. Note that the evaluation metrics (e.g., mIoU, chamfer distance) are also calculated during evaluation.
Visualizations of our complete point cloud and its reconstructed meshes are shown below. Please note that in our point cloud visualization, each color represents a specific part label within a category. For example, in the airplane
category, red indicates the plane body, green for the wings, and yellow for the engine.
Category | Partial Point Cloud | Complete Point Cloud (Ours) | Reconstructed Meshes |
---|---|---|---|
airplane |
![]() |
![]() |
![]() |
car |
![]() |
![]() |
![]() |
chair |
![]() |
![]() |
![]() |
table |
![]() |
![]() |
![]() |
@inproceedings{huang2022spovt,
title = {SPoVT: Semantic-Prototype Variational Transformer for Dense Point Cloud Semantic Completion},
author = {Huang, Sheng Yu and Hsu, Hao-Yu and Wang, Frank},
booktitle = {Advances in Neural Information Processing Systems},
pages = {33934--33946},
year = {2022},
}
This work is supported in part by the Tron Future Tech Inc. and National Science and Technology Council via NSTC-110-2634-F-002-052. We also thank National Center for High-performance Computing (NCHC) for providing computational and storage resources.