Golriz-code / PVT

PVT: Point-Voxel Transformer for 3D Deep Learning

PVT: Point-Voxel Transformer for 3D Deep Learning

Paper and Citation

The paper can be downloaded from arXiv.
If you like our work and think it helpful to your project, please cite it as follows.

@article{zhang2021point,
  title={PVT: Point-Voxel Transformer for 3D Deep Learning},
  author={Zhang, Cheng and Wan, Haocheng and Liu, Shengqiang and Shen, Xinyi and Wu, Zizhao},
  journal={arXiv preprint arXiv:2108.06076},
  year={2021}
}

Prerequisites

The code is built with following libraries (see requirements.txt):

• Python >= 3.7
• PyTorch >= 1.3
• numba
• numpy
• scipy
• six
• tensorboardX >= 1.2
• plyfile
• h5py
• timm

Data Preparation

ModelNet40

Download alignment ModelNet40 here and save in data/modelnet40_normal_resampled/.

S3DIS

We follow the data pre-processing in PointCNN. The code for preprocessing the S3DIS dataset is located in data/s3dis/. One should first download the dataset from here, then run

python data/s3dis/prepare_data.py

ShapeNet

We follow the data pre-processing in PointNet2. Please download the dataset from hereand save in data/shapenetcore_partanno_segmentation_benchmark_v0_normal/.

KITTI

For Frustum-PointNet backbone, we follow the data pre-processing in Frustum-Pointnets. One should first download the ground truth labels from here, then run

unzip data_object_label_2.zip
mv training/label_2 data/kitti/ground_truth
./data/kitti/frustum/download.sh

Pretrained Models

Here we provide a pretrained model on ModelNet40. The accuracy might vary a little bit compared to the paper, since we re-train some of the models for reproducibility. The path of the model is in ./checkpoints/cls/model.t7

Example training and testing

#train
python main_cls.py --exp_name=cls --num_points=1024 --use_sgd=True --batch_size 32 --epochs 200 --lr 0.001

#test
python main_cls.py --exp_name=cls --num_points=1024 --use_sgd=True --eval=True --model_path=checkpoints/cls/model.t7 --test_batch_size 32