/polarmix_scaling

PyTorch implementation of the NeurIPS2022 paper "PolarMix: A General Data Augmentation Technique for LiDAR Point Clouds"

Primary LanguagePythonMIT LicenseMIT

PolarMix

PolarMix: A General Data Augmentation Technique for LiDAR Point Clouds

Aoran Xiao, Jiaxing Huang, Dayan Guan, Kaiwen Cui, Shijian Lu, Ling Shao

NeurIPS 2022

News

[2023-04] PolarMix has been integrated into the MMDetection3D codebase. Check its new feature.
[2022-09-20] Code is released.
[2022-09-15] Our paper is accepted to NeurIPS 2022.

Usage

Installation

Please visit and follow installation instruction in this repo.

Data Preparation

SemanticKITTI

  • Please follow the instructions from here to download the SemanticKITTI dataset (both KITTI Odometry dataset and SemanticKITTI labels) and extract all the files in the sequences folder to /dataset/semantic-kitti. You shall see 22 folders 00, 01, …, 21; each with subfolders named velodyne and labels.
  • Change the data root path in configs/semantic_kitti/default.yaml

Training

SemanticKITTI

We release the training code for SPVCNN and MinkowskiNet with PolarMix. You may run the following code to train the model from scratch.

SPVCNN:

python train.py configs/semantic_kitti/spvcnn/cr0p5.yaml --run-dir runs/semantickitti/spvcnn_polarmix --distributed False

MinkowskiNet:

python train.py configs/semantic_kitti/minkunet/cr0p5.yaml --run-dir run/semantickitti/minkunet_polarmix --distributed False
  • Note we only used one 2080Ti for training and testing. Training from scratch takes around 1.5 days. You may try larger batch size or distributed learning for faster training.

Testing Models

You can run the following command to test the performance of SPVCNN/MinkUNet models with PolarMix.

torchpack dist-run -np 1 python test.py --name ./runs/semantickitti/spvcnn_polarmix
torchpack dist-run -np 1 python test.py --name ./runs/semantickitti/minkunet_polarmix

We provide pre-trained models of MinkUNet and SPVCNN. You may download and place them under './runs/semantickitti/' for testing.

mIoUs over validation set of SemanticKITTI are reported as follows:

w/o PolarMix w/ PolarMix
MinkUNet 58.9 65.0
SPVCNN 60.7 66.2

Visualizations

Follow instructions in this repo.

Citation

If you use this code for your research, please cite our paper.

@article{xiao2022polarmix,
  title={PolarMix: A General Data Augmentation Technique for LiDAR Point Clouds},
  author={Xiao, Aoran and Huang, Jiaxing and Guan, Dayan and Cui, Kaiwen and Lu, Shijian and Shao, Ling},
  journal={arXiv preprint arXiv:2208.00223},
  year={2022}
}

Thanks

We thank the opensource project TorchSparse and SPVNAS.

Related Repos

Find our other repos for point cloud understanding!