基于点云分割网络RandLA-Net的端到端部署

准备工作

1. 基于开源代码仓库地址：https://github.com/liuxuexun/RandLA-Net-Pytorch-New.git。
2. 已经基于该库复现了s3dis公共数据集或者自己的s3dis格式数据集，并导出了模型checkpoint.tar格式。自定义数据集标注，制作成s3dis格式以及进行完整训练将在下一篇文章进行详细介绍。
3. 熟悉RandLA-Net训练策略。

完整推理过程

准备：

# 换成你训练的类别
gt_class2label = {
   0: 'back', 1: 'box'} 
label_values = np.sort([k for k, v in gt_class2label.items()])
# 训练好的模型
model_path = r'./checkpoint.tar'
# 推理点云文件，即在网络训练时输入的点云格式
point_path = r'./Scene_1.txt'
# 推理前准备操作
device = torch.device('cuda:0')
net = Network(cfg)
optimizer = Adam(net.parameters(), lr=cfg.learning_rate)
checkpoint = torch.load(model_path, map_location='cpu')
net.load_state_dict(checkpoint['model_state_dict'])
net.to(device)
net.eval()
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

数据载入：

pc = np.loadtxt(r"./Scene_1.txt")
xyz_min = np.min(pc, axis=0)[0:3]
xyz, colors = (pc[:, :3] - xyz_min).astype(np.float32), pc[:, 3:6].astype(np.float32)
# 我的数据grid_sub_sampling操作前后数据一样，windows暂且使用该方式
sub_xyz, sub_colors = xyz, colors
# linux 使用该方式
# sub_xyz, sub_colors = DP.grid_sub_sampling(points=xyz, features=colors, grid_size=cfg.sub_grid_size)

# 替换KDTree使用cKDTree加速
search_tree = cKDTree(sub_xyz)
sub_colors = sub_colors / 255.0
input_trees = search_tree
input_colors = sub_colors
input_names = 'infer'
proj_idx = search_tree.query(xyz, k=1)[1].squeeze()
# 获得每个点云的最近邻索引
val_proj = proj_idx.astype(np.int32)

# 初始化每个点云属于每个类别的概率
test_probs = [np.zeros(shape=[input_colors.shape[0], len(gt_class2label)], dtype=np.float32)]

预处理过程：

# 基本同仓库代码一样
def tf_map(batch_xyz, batch_features, batch_pc_idx, batch_cloud_idx