RandLA-Net训练自己的数据集（基于S3DIS)

本次实验所用的是服务器上的GPU：RTX 2080 Ti(11GB) ，环境配置为：Ubuntu18.04、TensorFlow1.15.0以及cuda 10.0，python3.6.0。之前就在该网络上跑通了S3DIS数据集，所以沿用之前的环境配置进行了数据拷贝以及实例迁移。出于需要对自己的古建筑数据集在该网络进行测评。所以本内容着重基于讲解如何修改参数和路径以便跑通网络，对于环境配置内容不重点阐述。内容按照py文件进行修改。

1.数据预处理

1.克隆存储库：

git clone --depth=1 https://github.com/QingyongHu/RandLA-Net && cd RandLA-Net

然后放在根目录下进行解压：unzip RandLA-Net-master.zip

2.设置python环境：

conda create -n randlanet python=3.5
source activate randlanet
pip install -r helper_requirements.txt
sh compile_op.sh

3.准备训练数据：

准备好自己的数据集，将其进行解压，在RandLA-Net-master目录下创建一个文件夹命名为“data"然后将数据放在该文件夹下。

2.data_prepare_s3dis.py

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

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

dataset_path = '/root/autodl-tmp/RandLA-Net-master/data/S3DIS' #修改数据路径
anno_paths = [line.rstrip() for line in open(join(BASE_DIR, 'meta/anno_paths.txt'))]
anno_paths = [join(dataset_path, p) for p in anno_paths]

gt_class = [x.rstrip() for x in open(join(BASE_DIR, 'meta/class_names.txt'))]
gt_class2label = {cls: i for i, cls in enumerate(gt_class)}

sub_grid_size = 0.04
original_pc_folder = join(dirname(dataset_path), 'original_ply')
sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(sub_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
out_format = '.ply'


def convert_pc2ply(anno_path, save_path):
    """
    Convert original dataset files to ply file (each line is XYZRGBL).
    We aggregated all the points from each instance in the room.
    :param anno_path: path to annotations. e.g. Area_1/office_2/Annotations/
    :param save_path: path to save original point clouds (each line is XYZRGBL)
    :return: None
    """
    data_list = []

    for f in glob.glob(join(anno_path, '*.txt')):
        class_name = os.path.basename(f).split('_')[0]
        if class_name not in gt_class:  # note: in some room there is 'staris' class..
            class_name = 'clutter'
        pc = pd.read_csv(f, header=None, delim_whitespace=True).values
        labels = np.ones((pc.shape[0], 1)) * gt_class2label[class_name]
        data_list.append(np.concatenate([pc, labels], 1))  # Nx7

    pc_label = np.concatenate(data_list, 0)
    xyz_min = np.amin(pc_label, axis=0)[0:3]
    pc_label[:, 0:3] -= xyz_min

    xyz = pc_label[:, :3].astype(np.float32)
    colors = pc_label[:, 3:6].astype(np.uint8)
    labels = pc_label[:, 6].astype(np.uint8)
    write_ply(save_path, (xyz, colors, labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class'])

    # save sub_cloud and KDTree file
    sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(xyz, colors, labels, sub_grid_size)
    sub_colors = sub_colors / 255.0
    sub_ply_file = join(sub_pc_folder, save_path.split('/')[-1][:-4] + '.ply')
    write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class'])

    search_tree = KDTree(sub_xyz)
    kd_tree_file = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_KDTree.pkl')
    with open(kd_tree_file, 'wb') as f:
        pickle.dump(search_tree, f)

    proj_idx = np.squeeze(search_tree.query(xyz, return_distance=False))
    proj_idx = proj_idx.astype(np.int32)
    proj_save = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_proj.pkl')
    with open(proj_save, 'wb') as f:
        pickle.dump([proj_idx, labels], f)


if __name__ == '__main__':
    # Note: there is an extra character in the v1.2 data in Area_5/hallway_6. It's fixed manually.
    for annotation_path in anno_paths:
        print(annotation_path)
        elements = str(annotation_path).split('/')
        out_file_name = elements[-3] + '_' + elements[-2] + out_format
        convert_pc2ply(annotation_path, join(original_pc_folder, out_file_name))

然后运行文件：

python utils/data_prepare_s3dis.py

3.helpler_tool.py

class ConfigS3DIS:
    k_n = 16  # KNN
    num_layers = 5  # Number of layers
    num_points =40960   #Number of input points,这个不用修改
    num_classes = 10 # Number of valid classes，修改为你自己数据集的类别
    sub_grid_size = 0.04  # preprocess_parameter

    batch_size =2 #6  # batch_size during training，这个也需要修改根据自己数据集情况，我的报错了修改后解决了
    val_batch_size =1 #20  # batch_size during validation and test，这个参数和上面一样按照自己情况修改
    train_steps = 500  # Number of steps per epochs
    val_steps = 100  # Number of validation steps per epoch

    sub_sampling_ratio = [4, 4, 4, 4, 2]  # sampling ratio of random sampling at each layer
    d_out = [16, 64, 128, 256, 512]  # feature dimension

    noise_init = 3.5  # noise initial parameter
    max_epoch = 100  # maximum epoch during training
    learning_rate = 1e-2  # initial learning rate
    lr_decays = {i: 0.95 for i in range(0, 500)}  # decay rate of learning rate

    train_sum_dir = 'train_log'
    saving = True
    saving_path =None

下面这个一定要修改，不然会导致权重维度不匹配，我就是这个内容忽略了导致困扰了好久。 

  @staticmethod
    def get_class_weights(dataset_name):
        # pre-calculate the number of points in each category
        num_per_class = []
        if dataset_name is 'S3DIS':
            '''num_per_class = np.array([3370714, 2856755, 4919229, 318158, 375640, 478001, 974733,
                                      650464, 791496, 88727, 1284130, 229758, 2272837], dtype=np.int32)'''
            num_per_class = np.array([3370714, 2856755, 4919229, 318158, 375640, 478001, 974733,
                                      650464, 791496, 88727], dtype=np.int32)#自己数据集类别对应其中的个数，数目不需要修改
        elif dataset_name is 'Semantic3D':
            num_per_class = np.array([5181602, 5012952, 6830086, 1311528, 10476365, 946982, 334860, 269353],
                                     dtype=np.int32)
        elif dataset_name is 'SemanticKITTI':
            num_per_class = np.array([55437630, 320797, 541736, 2578735, 3274484, 552662, 184064, 78858,
                                      240942562, 17294618, 170599734, 6369672, 230413074, 101130274, 476491114,
                                      9833174, 129609852, 4506626, 1168181])
        weight = num_per_class / float(sum(num_per_class))
        ce_label_weight = 1 / (weight + 0.02)
        return np.expand_dims(ce_label_weight, axis=0)

4.main_S3DIS.py

class S3DIS:
    def __init__(self, test_area_idx):
        self.name = 'S3DIS'
        self.path = '/root/autodl-tmp/RandLA-Net-master/data/S3DIS'#路径修改为数据集所在位置
        self.label_to_names = {0: 'fang',
                               1: 'door',
                               2: 'floor',
                               3: 'roof',
                               4: 'stair',
                               5: 'taiji',
                               6: 'wall',
                               7: 'window',
                               8: 'zhuzi',
                               9: 'clutter'}#修改为自己数据集类别
        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.array([])

        self.val_split = 'Area_' + str(test_area_idx)
        self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply'))

        # Initiate containers
        self.val_proj = []
        self.val_labels = []
        self.possibility = {}
        self.min_possibility = {}
        self.input_trees = {'training': [], 'validation': []}
        self.input_colors = {'training': [], 'validation': []}
        self.input_labels = {'training': [], 'validation': []}
        self.input_names = {'training': [], 'validation': []}
        self.load_sub_sampled_clouds(cfg.sub_grid_size)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]')
    parser.add_argument('--test_area', type=int, default=3, help='Which area to use for test, option: 1-6 [default: 5]')#改为自己测试的区域
    parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis')
    parser.add_argument('--model_path', type=str, default='None', help='pretrained model path')
    FLAGS = parser.parse_args()

 5. 6_fold_cv.py

if __name__ == '__main__':
    base_dir = '/root/autodl-tmp/RandLA-Net-master/data/S3DIS/results'#结果ply文件存放路径
    original_data_dir = '/root/autodl-tmp/RandLA-Net-master/data/S3DIS/original_ply'#源数据存放路劲
    data_path = glob.glob(os.path.join(base_dir, '*.ply'))
    data_path = np.sort(data_path)

    test_total_correct = 0
    test_total_seen = 0
    gt_classes = [0 for _ in range(10)]
    positive_classes = [0 for _ in range(10)]
    true_positive_classes = [0 for _ in range(10)]#将13改为10，因为我的类别为10
    visualization =True#将False改为True用于结果可视化

    iou_list = []
    for n in range(10):
        iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n])
        iou_list.append(iou)
    mean_iou = sum(iou_list) /10.0
    print('eval accuracy: {}'.format(test_total_correct / float(test_total_seen)))
    print('mean IOU:{}'.format(mean_iou))
    print(iou_list)

    acc_list = []
    for n in range(10):
        acc = true_positive_classes[n] / float(gt_classes[n])
        acc_list.append(acc)
    mean_acc = sum(acc_list) / 10.0
    print('mAcc value is :{}'.format(mean_acc))
将以上10改为你1自己的类别数目

 6.jobs_6_fold_cv_s3dis.sh

我的区域只有三个所以六折交叉验证将后三个注释掉：

python -B main_S3DIS.py --gpu 0 --mode train --test_area 1
python -B main_S3DIS.py --gpu 0 --mode test --test_area 1
python -B main_S3DIS.py --gpu 0 --mode train --test_area 2
python -B main_S3DIS.py --gpu 0 --mode test --test_area 2
python -B main_S3DIS.py --gpu 0 --mode train --test_area 3
python -B main_S3DIS.py --gpu 0 --mode test --test_area 3
# python -B main_S3DIS.py --gpu 0 --mode train --test_area 4
# python -B main_S3DIS.py --gpu 0 --mode test --test_area 4
# python -B main_S3DIS.py --gpu 0 --mode train --test_area 5
# python -B main_S3DIS.py --gpu 0 --mode test --test_area 5
# python -B main_S3DIS.py --gpu 0 --mode train --test_area 6
# python -B main_S3DIS.py --gpu 0 --mode test --test_area 6

修改以上内容就可以跑通自己的数据集了，注意这只是基本情况具体自己的数据集还需要结合实际，本人也在学习当中有许多内容不足可以互相讨论。