点云语义分割与部署项目

本仓库包含点云语义分割模型的训练、测试和部署全流程代码，支持多种模型架构和分割任务，包括多类和二分类任务，并实现了知识蒸馏技术优化模型性能。

项目特点

• 支持DGCNN、PointNet++等多种点云处理模型
• 实现知识蒸馏技术优化小模型性能
• 提供完整的训练、测试和部署流程
• 针对圆形电厂和条形结构等特殊场景优化

主要脚本功能

训练脚本

1. train_semseg_rs.py (对应 a.py)

• 描述: 训练动态图卷积神经网络 (DGCNN) 模型用于点云语义分割。
• 任务: 多类分割，包含 5 个类别（liaodui（料堆）、dabi（大臂）、guidao（轨道）、madao（马道）、zaodian（噪点））。
• 特点: 使用加权交叉熵损失，支持数据增强（绕 Z 轴随机旋转），每 5 个 epoch 保存一次模型，并根据均交并比 (mIoU) 保存最佳模型。
• 数据路径: ./mydata/data_npy/
• 测试区域: 2
• 默认训练轮数: 100

2. train_semseg_rshgnn.py (对应 b.py)

• 描述: 训练 DGCNN 的层次图神经网络变体 (DGCNN-HGNN) 用于点云语义分割。
• 任务: 多类分割，包含 5 个类别（与 train_semseg_rs.py 相同）。
• 特点: 类似于 train_semseg_rs.py，但使用更先进的 dgcnnhgnn 模型，采用加权交叉熵损失并定期保存模型。
• 数据路径: ./mydata/data_npy/
• 测试区域: 2
• 默认训练轮数: 100

3. train_semseg-kd.py (对应 c.py)

• 描述: 训练 PointNet++ 学生模型，通过 DGCNN 教师模型的知识蒸馏进行点云语义分割。
• 任务: 二分类（detect（检测到）、no_detect（未检测到））。
• 特点: 结合加权交叉熵损失和 KL 散度损失（可选注意力损失，目前注释掉），训练结束后保存最终模型，因蒸馏效率使用较少轮数。
• 数据路径: ./data_npy/
• 测试区域: 5
• 默认训练轮数: 50
• 额外参数: temp（蒸馏温度）、alpha（蒸馏损失权重）、beta（注意力损失权重）。

4. train_semseg-pointnet2.py (对应 d.py)

• 描述: 训练 PointNet++ 模型用于点云语义分割。
• 任务: 二分类（detect（检测到）、no_detect（未检测到））。
• 特点: 使用加权交叉熵损失，支持数据增强，每 5 个 epoch 保存一次模型，并根据 mIoU 保存最佳模型。
• 数据路径: ./data_npy/
• 测试区域: 5
• 默认训练轮数: 100

5. train_semseg.py (对应 e.py)

• 描述: 训练 DGCNN 模型用于点云语义分割。
• 任务: 二分类（detect（检测到）、no_detect（未检测到））。
• 特点: 使用加权交叉熵损失，支持数据增强，每 5 个 epoch 保存一次模型，并根据 mIoU 保存最佳模型。
• 数据路径: ./data_npy/
• 测试区域: 5
• 默认训练轮数: 100

测试脚本

1. test_semseg_cir.py

• 描述: 用于圆形电厂点云数据的测试
• 输入: 预训练模型(.pth文件)和测试数据
• 输出: 分割结果和评估指标
• 可视化: 生成圆形电厂分割效果图(cir1.jpg, cir2.jpg)

2. test_pred_trained.py

• 描述: 通用点云分割测试脚本
• 功能: 加载训练好的模型进行推理测试
• 支持: 多种点云数据格式输入

3. test_scripting.py

• 描述: 模型导出和脚本化测试
• 功能: 将PyTorch模型导出为TorchScript格式用于部署

部署说明

项目已提供部署所需的模型文件:

• deploy/dgcnn_segmentationcir.pt: 针对圆形电厂优化的部署模型

部署流程:

1. 使用test_scripting.py将训练好的模型导出为TorchScript格式
2. 通过deploy/目录下的模型文件进行推理
3. 集成到目标部署环境(C++/Python均可)

效果展示

• 圆形电厂测试效果:
• 条形结构测试效果:

依赖项

• Python 3.x
• PyTorch: 确保安装了支持 GPU 的 PyTorch（如果使用 CUDA）。
• 其他依赖: numpy, torch, pathlib, tqdm, logging，以及模型特定的工具（例如 provider, OrangeDataset）。

安装依赖项：

pip install torch numpy pathlib tqdm open3d matplotlib

使用建议

1. 训练时可根据显存大小调整batch_size参数
2. 部署前建议使用test_scripting.py验证模型导出正确性
3. 圆形电厂场景推荐使用专门的dgcnn_segmentationcir.pt模型