PointNet++训练自定义数据集

2024-10-28 17:39:48

我们已经学习了PointNet++的网络结构，以及在S3Dis数据集上的训练、测试以及可视化等操作，那么如何使用该模型来训练我们自己的数据集呢，其实，操作十分简单，只需要修改几个参数即可。

数据集查看

首先，我们需要获得数据集，这里我们的数据集已经经过处理为npy格式，其维度与S3Dis数据集的维度相同，不同之处在于，该数据集中没有rgb维度，但为了保持一致，这里我们将其填充0。

import numpy as np
import open3d as o3d
data=np.load("datasets/Area_1_tower_1.npy")
pcd=o3d.geometry.PointCloud()
pcd.points=o3d.utility.Vector3dVector(data[:,0:3])
o3d.visualization.draw_geometries([pcd])

模型训练

随后，我们修改一下相关参数即可。
首先是train_semseg.py文件中，修改类别，当然我们还需要根据硬件配置来修改batch等信息

classes = [
		"tower",
		"wire",
		"ground_wire",
		"insulator",
		"drainage_thread",
		"tower_distributing",
		"wire_distributing",
		"plant",
		"road",
 ]

然后是数据集的地址与类别数量：

root = 'data/datasets/'
NUM_CLASSES = 9

随后便是修改加载数据集的文件S3DISDataLoader.py，在该数据集中主要还是修改与类别相关的参数，因为由于不同类别的点云数量不同，因此在求损失时需要进行均衡化。

labelweights = np.zeros(9)#原本为13，改为9
tmp, _ = np.histogram(labels, range(10))#原本为14改为10

此外，在读取数据集时，采用的是指定测试集编号，其余的算做训练集的方式，因此这里我们虽然只有一个数据集，但需要将其复制一份，并起不同的名称：

if split == 'train':
      rooms_split = [room for room in rooms if not 'Area_{}'.format(test_area) in room]
else:
      rooms_split = [room for room in rooms if 'Area_{}'.format(test_area) in room]

我们指定的测试集编号为2。