PCDet: 3D Point Cloud Detection
Introduction
PCDet是一种用于点云3D对象感知的基于pytorch的代码。它支持多种当前优秀的三维点云感知算法(PointPillar, SECOND, Part-A^2 Net),在第一或第二阶段框架提供高度的重构。
同时也是Part-A^2 Net的官方代码。
注意:当前框架主要是包括基于点素的方法,我们将继续支持更多基于点云的方法。
Currently Supported Features
- 同时支持一阶段和二阶段的三维物体感知框架
- 在KITTI数据集上通过分布式多GPU训练5小时获得很好结果
- 清晰的代码结构支持更多数据集和方法
- Rol-aware点云集中
- GPU版本的3D IoU calculation和rotated NMS
Model Zoo
支持方法在下表中,这些结果是在kitty数据集car物体的验证集上的结果,所有的模型都是在8GPUs上训练的,模型提供如下:
\ | training time | AP_Easy | AP_Mod | AP_Hard | download |
---|---|---|---|---|---|
PointPillar | ~2hours | 87.37 | 77.30 | 74.02 | model-18M |
SECOND | ~2hours | 88.46 | 78.46 | 76.63 | model-20M |
Part-A^2 | ~5hours | 89.66 | 79.45 | 78.80 | model-209M |
Part-A^2-fc | ~5hours | 89.57 | 79.31 | 78.61 | model-244M |
Installation
Requirements
代码测试环境如下:
- Linux (tested on Ubuntu 14.04/16.04)
- Python 3.6+
- PyTorch 1.1 or higher
- CUDA 9.0 or higher
Install pcdet
1、拷贝本工程
git clone https://github.com/sshaoshuai/PCDet.git
2、安装依赖包
pip install -r requirements.txt
3、安装pcdet库
python setup.py develop
Dataset Preparation
当前只支持KITTI数据集,欢迎贡献更多的数据集支持
KITTI Dataset
请下载官方KITTY 3D对象检测的数据集,并且按如下方式组织数据集(道路数据集在此下载,虽然在训练中这是可选数据):
PCDet
├── data
│ ├── kitti
│ │ │──ImageSets
│ │ │──training
│ │ │ ├──calib & velodyne & label_2 & image_2 & (optional: planes)
│ │ │──testing
│ │ │ ├──calib & velodyne & image_2
├── pcdet
├── tools
通过如下代码产生数据信息:
python kitti_dataset.py create_kitti_infos
Getting Started
所有的配置信息都在tools/cfgs/文件下
Test and evaluate the pretrained models
- 测试已训练好的模型:
python test.py --cfg_file ${CONFIG_FILE} --batch_size 4 --ckpt ${CKPT}
- 例如,测试上面提供的Part-A^2模型,运行如下代码(此处我们添加-set 去修改一些默认参数,来匹配训练好的模型,其他提供的训练好的模型不需要加):
python test.py --cfg_file cfgs/PartA2_car.yaml --batch_size 4 --ckpt PartA2.pth \
--set MODEL.RPN.BACKBONE.NAME UNetV0 MODEL.RPN.RPN_HEAD.ARGS use_binary_dir_classifier:True
- 为了估计一个训练集的所有检测点并在tensorboard上画出表现图,增加 --eval_all参数:
python test.py --cfg_file ${CONFIG_FILE} --batch_size 4 --eval_all
Train a model
当前,为了训练PointPillar、SECOND 、PartA2,batch_size取决于你训练GPU的数量,我们使用了 4*gpu数量的batch_size,也就是32。
- 多gpus训练
sh scripts/dist_train.sh ${NUM_GPUS} \
--cfg_file ${CONFIG_FILE} --batch_size ${BATCH_SIZE}
- 多机器训练
sh scripts/slurm_train.sh ${PARTITION} ${JOB_NAME} ${NUM_GPUS} \
--cfg_file ${CONFIG_FILE} --batch_size ${BATCH_SIZE}
- 单gpu训练
python train.py --cfg_file ${CONFIG_FILE} --batch_size ${BATCH_SIZE}
Acknowledgement
我们感谢second.pytorch提供了一阶段体素框架,且借鉴了PointRCN的代码。
我们希望这个工程能够足够的健壮和灵活,来加快以前的工作,或者产生出新的方法。
Citation
如果你发现此工作对你的工作有帮助,请如下标注:
@article{shi2020points,
title={From Points to Parts: 3D Object Detection from Point Cloud with Part-aware and Part-aggregation Network},
author={Shi, Shaoshuai and Wang, Zhe and Shi, Jianping and Wang, Xiaogang and Li, Hongsheng},
journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
year={2020},
publisher={IEEE}
}
或者
@inproceedings{shi2019pointrcnn,
title={PointRCNN: 3d Object Progposal Generation and Detection from Point Cloud},
author={Shi, Shaoshuai and Wang, Xiaogang and Li, Hongsheng},
booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
pages={770--779},
year={2019}
}
Contact
Should you have any question, please contact Shaoshuai Shi (@sshaoshuai).