【代码阅读】RandLA-Net

RandLa-Net是新提出来的针对大场景语义分割的方法，效果拔群。我对该文章的解读可以看我另外 一篇博客，论文作者给出的代码是 TensorFlow版本的。

接下来，我们就看一下这个代码都做了什么操作。以SemanticKitti为例。

由于作者给出的代码是在TensorFlow下的，我改到了Pytorch下面，代码详见我的Github，是实现了SemanticKITTI数据集下的训练。

数据预处理

# utils/data_prepare-semantickitti.py

# line 42-50
points = DP.load_pc_kitti(join(pc_path, scan_id))
labels = DP.load_label_kitti(join(label_path, str(scan_id[:-4]) + '.label'), remap_lut)
sub_points, sub_labels = DP.grid_sub_sampling(points, labels=labels, grid_size=grid_size)
search_tree = KDTree(sub_points)
KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl')
np.save(join(pc_path_out, scan_id)[:-4], sub_points)
np.save(join(label_path_out, scan_id)[:-4], sub_labels)
with open(KDTree_save, 'wb') as f:
    pickle.dump(search_tree, f)

可以看到，上述预处理，是把point和label做了grid sampling，并且生成了一个kdtree保存下来。

dataset

是用main_SemanticKITTI/SemanticKITTI这个类实现的，我们看看这部分做了什么

class SemanticKITTI:
    def __init__(self, test_id):
        ...
        
    # Generate the input data flow
    def get_batch_gen(self, split):
        ...
        def spatially_regular_gen():
            # Generator loop
            # line 72-79
            for i in range(num_per_epoch):
                if split != 'test':
                    cloud_ind = i
                    pc_path = path_list[cloud_ind]
                    pc, tree, labels = self.get_data(pc_path)
                    # crop a small point cloud
                    pick_idx = np.random.choice(len(pc), 1)
                    selected_pc, selected_labels, selected_idx = self.crop_pc(pc, labels, tree, pick_idx)
                ...
        ...
        return gen_func, gen_types, gen_shapes

    def get_data(self, file_path):  #从file_path所指向的文件中读入point，kdtree和label
        ...
        return points, search_tree, labels

    @staticmethod
    def crop_pc(points, labels, search_tree, pick_idx):
        # crop a fixed size point cloud for training
        center_point = points[pick_idx, :].reshape(1, -1)
        select_idx = search_tree.query(center_point, k=cfg.num_points)[1][0]
        select_idx = DP.shuffle_idx(select_idx)
        select_points = points[select_idx]
        select_labels = labels[select_idx]
        return select_points, select_labels, select_idx

    @staticmethod
    def