复现RandLA-Net网络，并训练自己的数据集

具体的复现过程可参考前面发的博客。

在训练自己的数据之前，首先要针对数据预处理进行修改，因为官方提供的数据中，标签是以.label格式单独储存于独立文件中，而我在处理标签是直接作为一列储存在点云文件中。此外，由于官方数据时具有颜色的，而我的数据是未着色的，因此直接将颜色赋0.

下面是修改后的data_prepare文件：

from sklearn.neighbors import KDTree
from os.path import join, exists, dirname, abspath
import numpy as np
import os, glob, pickle
import sys

BASE_DIR = dirname(abspath(__file__))
ROOT_DIR = dirname(BASE_DIR)
sys.path.append(BASE_DIR)
sys.path.append(ROOT_DIR)
from helper_ply import write_ply
from helper_tool import DataProcessing as DP

grid_size = 0.06  #我的点云数据集比较密，所以下采样间隔取大一点
dataset_path = 'data/semantic_kitti/out_16'
original_pc_folder = join(dirname(dataset_path), 'original_ply')
sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(grid_size))
os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None
os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None

for pc_path in glob.glob(join(dataset_path, '*.txt')):
    print(pc_path)
    # file_name = pc_path.split('/')[-1][:-4]
    file_name=os.path.basename(pc_path)[:-4]

    # check if it has already calculated
    if exists(join(sub_pc_folder, file_name + '_KDTree.pkl')):
        continue

    #直接使用numpy加载txt格式的点云文件
    pc=np.loadtxt(pc_path)

    #我的数据集中训练和测试集都是x,y,z,r,g,b,label的存储格式
    labels=pc[:,-1].astype(np.uint8)

    full_ply_path = join(original_pc_folder, file_name + '.ply')

    #  Subsample to save space
    sub_points, sub_colors, sub_labels = DP.grid_sub_sampling(pc[:, :3].astype(np.float32),
                                                              pc[:, 3:6].astype(np.uint8), labels, 0.01)
    sub_labels = np.squeeze(sub_labels)

    write_ply(full_ply_path, (sub_points, sub_colors, sub_labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class'])

    # save sub_cloud and KDTree file
    sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(sub_points, sub_colors, sub_labels, grid_size)
    sub_colors = sub_colors / 255.0
    sub_labels = np.squeeze(sub_labels)
    sub_ply_file = join(sub_pc_folder, file_name + '.ply')
    write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class'])

    search_tree = KDTree(sub_xyz, leaf_size=50)
    kd_tree_file = join(sub_pc_folder, file_name + '_KDTree.pkl')
    with open(kd_tree_file, 'wb') as f:
        pickle.dump(search_tree, f)

    proj_idx = np.squeeze(search_tree.query(sub_points, return_distance=False))
    proj_idx = proj_idx.astype(np.int32)
    proj_save = join(sub_pc_folder, file_name + '_proj.pkl')
    with open(proj_save, 'wb') as f:
        pickle.dump([proj_idx, labels], f)

此外，main文件中的类定义也要修改，将分类标签改成自己数据所对应的，同时要将训练集和测试集的划分代码进行修改。

    def __init__(self):
        self.name = 'power'
        self.path = 'utils/data/semantic_kitti'
        self.label_to_names = {0: 'background',
                               1: 'build',
                               }
        self.num_classes = len(self.label_to_names)
        self.label_values = np.sort([k for k, v in self.label_to_names.items()])
        self.label_to_idx = {l: i for i, l in enumerate(self.label_values)}
        self.ignored_labels = np.sort([])

        self.original_folder = join(self.path, 'out_16')
        self.full_pc_folder = join(self.path, 'original_ply')
        self.sub_pc_folder = join(self.path, 'input_{:.3f}'.format(0.06))

        # 训练、验证、测试数据都在original_data数据集中，需要做划分


        # Initial training-validation-testing files
        self.train_files = []
        self.val_files = []
        self.test_files = []
        cloud_names = [file_name[:-4] for file_name in os.listdir(self.original_folder) if file_name[-4:] == '.txt']
        self.val_split = cloud_names[:30]  # "file1"是点云文件的名字
        self.test_split = cloud_names[30:45]
        # 根据文件名划分训练、验证、测试数据集
        for pc_name in cloud_names:
            pc_file = join(self.sub_pc_folder, pc_name + '.ply')
            if pc_name in self.val_split:
                self.val_files.append(pc_file)
            elif pc_name in self.test_split:
                self.test_files.append(pc_file)
            else:
                self.train_files.append(pc_file)

        # Initiate containers
        self.val_proj = []
        self.val_labels = []
        self.test_proj = []
        self.test_labels = []

        self.possibility = {}
        self.min_possibility = {}
        self.class_weight = {}
        self.input_trees = {'training': [], 'validation': [], 'test': []}
        self.input_colors = {'training': [], 'validation': [], 'test': []}
        self.input_labels = {'training': [], 'validation': []}

        self.load_sub_sampled_clouds(cfg.sub_grid_size)

此外，数据集的权重也要进行修改，因为我是进行单类别分割，因此只是将权值全部在设置为1.

最后，按照正常的训练流程进行训练即可。