在Google Colab中训练3D物体检测模型——VoteNet

Deep Hough Voting for 3D Object Detection in Point Clouds
作者：Charles R. Qi，Or Litany，何恺明，Leonidas J. Guibas等
简介：当前主流的3D目标检测方法，很大层度上受2D检测器的影响。充分利用2D检测器提取3D建议框（例如Frustum-PointNets），或将3D点云转换为规则的网格（如VoxelNet），VoteNet是一种直接处理原始数据、不依赖任何2D检测器的点云3D检测框架。在这篇论文中，研究人员回归第一原则，为点云数据构建了一个尽可能通用的3D检测pipeline。然而，由于数据的稀疏性，直接从场景点预测边界框参数时面临一个主要挑战：一个3D物体的质心可能远离任何表面点，因此很难用一个步骤准确地回归。为了解决这一问题，研究人员提出VoteNet，这是一个基于深度点集网络和霍夫投票的端到端3D目标检测网络。该模型设计简单，模型尺寸紧凑，而且效率高，在SUN RGB-D和ScanNet两大真实3D扫描数据集上实现了最先进的3D检测精度。值得注意的是，VoteNet优于以前的方法，而且不依赖彩色图像，使用纯几何信息。

使用Colab训练
一、挂载谷歌硬盘

from google.colab import drive 
drive.mount("/content/drive")

二、安装/降级依赖库

!pip install plyfile
!pip install trimesh==2.35.39
!pip install networkx==2.2

三、git clone代码

cd drive/My Drive
!git clone https://github.com/facebookresearch/votenet.git

四、运行demo
1、为Pointnet++编译CUDA层

cd pointnet2
python setup.py install

2、下载数据集
https://drive.google.com/file/d/1oem0w5y5pjo2whBhAqTtuaYuyBu1OG8l/view?usp=sharing（谷歌硬盘的分享链接，打开链接将文件转存到votenet文件夹），然后在colab中解压。

!unzip votenet_pretrained_models.zip 

3、运行demo.py

!python demo.py

该demo使用预先训练的模型，从含有桌子和几把椅子的室内空间点云（来自SUN RGB-D val set）中检测对象。 我们可以使用3D可视化软件（例如MeshLab）来查看保存在demo_files/sunrgbd_results中的结果。其中，打开***_pc.ply 和***_pred_confident_nms_bbox.ply 可以看到输入的原始点云和预测出来的 3D bounding boxes(如下两图)


五、训练模型
可以使用SUN RGB-D数据集训练，也可以使用ScanNet数据集训练，下面只介绍使用SUN RGB-D进行训练的过程：
1、数据集预处理
这一步的具体过程可以查看votenet/sunrgbd里面的README文件，因为需要用到Matlab处理，无法在云端运行，所以这一步我将必要的数据下载到自己电脑处理后再上传到谷歌硬盘。
2、训练

CUDA_VISIBLE_DEVICES=0 python train.py --dataset sunrgbd --log_dir log_sunrgbd

六、测试模型

python eval.py --dataset sunrgbd --checkpoint_path log_sunrgbd/checkpoint.tar --dump_dir eval_sunrgbd --cluster_sampling seed_fps --use_3d_nms --use_cls_nms --per_class_proposal

Reference:作品GitHub
论文地址
SUN RGB-D数据集介绍