【PointNet++】PointNet++复现（PyTorch版本）

【PointNet++】PointNet++复现（PyTorch版本）

0. 相关文章

PointNet模型搭建
基于自建数据训练PointNet分割网络

1. PointNet++代码下载

本文搭建的是PyTorch版本的PointNet++代码仓库地址为：PyTorch版本

补充：TensorFlow版本

2. 数据集下载

PointNet++有三种基础任务：对象分类、对象部件分割、场景语义分割，因此需要三个数据集进行测试

ModelNet40：用于对象分类任务的数据集，下载地址：点击此处，百度网盘链接：【链接: https://pan.baidu.com/s/1jeeuS6K4zGUSo-tDJpyxJA 提取码: ruxe 】

ModelNet40数据集内容包含47个对象类别及7个相关配置文件，如图所示：

ShapeNet：用于对象部件分割任务的数据集，下载地址：点击此处，百度网盘链接：【链接: https://pan.baidu.com/s/1UJBICzt4hHJtyUuB1m2lVg 提取码: uuyh 】

shapenetcore_partanno_segmentation_benchmark_v0_normal数据集内容包含16个对象数据、synsetoffset2category.txt数据配置文件和train_test_split数据划分配置文件，如图所示：

S3DIS：用于场景语义分割任务的数据集，下载地址：点击此处，百度网盘链接：【链接: https://pan.baidu.com/s/1Zx4KbyCgll0ejv-8cQ5Z7Q 提取码: yf54 】

S3DIS数据集包含6个场景数据，每个场景中又包含多种类别的点云数据，如图所示：

3. 环境介绍

Ubuntu 22.04，NVIDIA GeForce RTX 3090

Python3.8，Pytorch1.7，cuda11.0

conda install pytorch==1.7.0 torchvision==0.8.0 torchaudio==0.7.0 cudatoolkit=11.0 -c pytorch

使用Anaconda虚拟环境根据PointNet++readme文件运行训练程序，缺少第三方包，通过pip install方式下载

4. 训练

4.1. 对象分类任务

在项目根目录下新建data目录，将ModelNet40数据集放至该目录中，如图所示：

在进行对象分类任务时有多个可供选择的配置项，如下所示：

默认训练方式：

python train_classification.py --model pointnet2_cls_ssg --log_dir pointnet2_cls_ssg

添加法向量特征的训练方式：

python train_classification.py --model pointnet2_cls_ssg --use_normals --log_dir pointnet2_cls_ssg_normal

使用均匀采样的训练方式：

python train_classification.py --model pointnet2_cls_ssg --use_uniform_sample --log_dir pointnet2_cls_ssg_fps

4.2. 对象部件分割任务

将ShapeNet数据集放入项目根目录data文件夹下，具体格式如图所示：

训练方法：

python train_partseg.py --model pointnet2_part_seg_msg --normal --log_dir pointnet2_part_seg_msg

4.3. 场景语义分割任务

将S3DIS数据集放至放入项目根目录data/s3dis文件夹下，具体格式如图所示：

然后基于此数据集制作PointNet++的数据集，使用以下程序：

cd data_utils
python collect_indoor3d_data.py

生成的数据集同样会保存在data目录下，其中数据格式是npy格式，如图所示：

训练方法：

python train_semseg.py --model pointnet2_sem_seg --test_area 5 --log_dir pointnet2_sem_seg

5. 测试

对象分类任务测试方法：

python test_classification.py --log_dir pointnet2_cls_ssg
python test_classification.py --use_normals --log_dir pointnet2_cls_ssg_normal
python test_classification.py --use_uniform_sample --log_dir pointnet2_cls_ssg_fps

对象部件分割任务测试方法：

python test_partseg.py --normal --log_dir pointnet2_part_seg_msg

场景语义分割测试方法：

python test_semseg.py --log_dir pointnet2_sem_seg --test_area 5 --visual

场景语义分割测试方法运行结束后，会在log/sem_seg/pointnet2_sem_seg/visual目录中生成每个场景分割后的obj文件，可以对齐进行可视化，可视化效果如图所示：

可视化脚本如下所示：

import open3d as o3d
import numpy as np


def visual_obj(path):
    with open(path, "r") as obj_file:
        points = []
        colors = []

        for line in obj_file.readlines():
            line = line.strip()
            line_list = line.split(" ")
            points.append(np.array(line_list[1:4]))
            colors.append(np.array(line_list[4:7]))

        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        pcd.colors = o3d.utility.Vector3dVector(colors)
        o3d.visualization.draw_geometries([pcd])


def main():
    objt_file_path = r"E:\07-code\Pointnet_Pointnet2_pytorch-master\test\Area_5_WC_1_pred.obj"
    visual_obj(objt_file_path)
    
    
if __name__ == '__main__':
    main()