关于S3DIS数据集预处理的代码解析

我用的是kpconvxgithub下说的一个点云分类检测框架下的预处理文件
叫做pointcept
github 官网如下
https://github.com/Pointcept/Pointcept

代码研读

主要针对这个文件进行研读，这个python是将txt转换成.npy格式

程序的主入口，首先是一个命令行参数，其中分别是数据集位置，raw位置，输出位置，和格式的选择，这个见github官网

这里确定一下说如果选择了法向量模式的话，确保raw路径不为空

这里的代码，就是将S3DIS数据集的每一个room 和 每一个area下的angle（即当前area下的全部点云数据）的名字放进列表中，其中的split就是我们在命令行中选择要进行数据转换的area

然后这一部分就是对raw部分进行处理，我没有实验，大概就是处理了一下然后放进了mesh中

然后就要开始对数据集进行处理了，这里传给processpoolexecutor的参数就是有几个在工作，命令行里的默认值是1，他还提示了说如果要处理法向量就要有110g的内存

2. ProcessPoolExecutor部分（对数据进行处理成npy）
   
   这里就是类的初始化，还是在设定max_worker，大致意思就是

• max_workers 为 None：

• self._max_workers = os.cpu_count() or 1：如果没有提供 max_workers 参数，则默认使用机器的 CPU 核心数来决定最大进程数（os.cpu_count() 返回当前机器的 CPU 核心数，如果 os.cpu_count() 返回 None，则使用 1 个进程）。

• 如果操作系统是 Windows (sys.platform == ‘win32’)，则限制最大进程数为 _MAX_WINDOWS_WORKERS，这可能是为了防止 Windows 上的进程过多导致不稳定。_MAX_WINDOWS_WORKERS 是一个预定义的常量，表示 Windows 系统支持的最大工作进程数。

• max_workers 不为 None：

• if max_workers <= 0:：如果指定的 max_workers 小于等于 0，抛出 ValueError 错误，因为进程池中的工作进程数必须大于 0。

• elif (sys.platform == ‘win32’ and max_workers > _MAX_WINDOWS_WORKERS):：如果是在 Windows 系统下，且 max_workers 超过了 Windows 系统允许的最大工作进程数，则抛出错误。
  最终，self._max_workers 将设置为指定的 max_workers。

这一部分基本都是关于进程的设置，具体意思为

• _queue_management_thread：用来管理队列的线程，目前初始化为 None，后续会用来管理进程池中队列的工作。
• _processes：存储进程池中各个进程的映射字典，以进程 ID (pid) 为键，进程对象为值。
• _shutdown_thread：用于标记是否已开始关闭进程池，初始化为 False。
• _shutdown_lock：用于在关闭进程池时保护资源，确保关闭操作是线程安全的。
• _broken：用于标记进程池是否已经“损坏”或出现故障。
• _queue_count：跟踪当前队列中的项目数。
• _pending_work_items：跟踪待处理的工作项。

这里的parse_room才是对其进行处理，而后面的六个参数才是parse_room的参数

3. pares_room阶段
   
   首先是classes的初始化和数据集路径和保存路径

这里的四个就是最终要生成的，其中语义标签是指他是桌子还是椅子，然后实例标签呢，是指他是第几个桌子还是第几个椅子，即每个点属于的类别

这里就开始处理了，其中object_path_list中包含全部的txt文件，

• 然后对于每个物体，比如一个table_1.txt，都读取出其中的txt中的全部点的坐标和color
• 然后classname如果 object_name（如 table）在 classes 里，就使用原名称；
  否则，设为 “clutter”，表示该物体是杂乱物体（即无类别的干扰物）
• 然后语义标签呢，classlabel【class_name】就是找到其类别对应的映射编码比如0,1，2,3然后复制n份，即一份txt有n个点，那么就复制n份
• 同理下面的实例标签也是一样，比如我这个文件名是table_1那么这个实例标签就是1，且复制了n份
  
  这个就是将各个的点云数据角度对齐，不仔细看了，直接选择就好

接下来就是将几个数组都垂直分布，分不成（n*3）的格式，np.ascontiguousarray(…):
确保生成的数组在内存中是 连续存储 的，提高计算效率，便于后续操作（如 GPU 加速）。

附上一个稍加修改可以直接用的

下面展示一些 内联代码片。

"""
Preprocessing Script for S3DIS
Parsing normal vectors has a large consumption of memory. Please reduce max_workers if memory is limited.

Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

import os
import argparse
import glob
import numpy as np

try:
    import open3d
except ImportError:
    import warnings

    warnings.warn("Please install open3d for parsing normal")

try:
    import trimesh
except ImportError:
    import warnings

    warnings.warn("Please install trimesh for parsing normal")

from concurrent.futures import ProcessPoolExecutor
from itertools import repeat

area_mesh_dict = {}


def parse_room(
    dataset_root, output_root, align_angle=False, parse_normal=False
):
    classes = [
        "ceiling",
        "floor",
        "wall",
        "beam",
        "column",
        "window",
        "door",
        "table",
        "chair",
        "sofa",
        "bookcase",
        "board",
        "clutter",
    ]
    class2label = {cls: i for i, cls in enumerate(classes)}
    #source_dir = os.path.join(dataset_root, room)
    source_dir = dataset_root
    #save_path = os.path.join(output_root, room)
    save_path = output_root
    print(save_path)
    print(source_dir)
    os.makedirs(save_path, exist_ok=True)
    object_path_list = sorted(glob.glob(os.path.join(source_dir, "Annotations/*.txt")))

    room_coords = []  #坐标信息
    room_colors = []  #颜色
    room_normals = []
    room_semantic_gt = []   #语义标签
    room_instance_gt = []   #实例标签

    for object_id, object_path in enumerate(object_path_list):
        object_name = os.path.basename(object_path).split("_")[0]
        obj = np.loadtxt(object_path)
        coords = obj[:, :3]
        colors = obj[:, 3:6]
        # note: in some room there is 'stairs' class
        class_name = object_name if object_name in classes else "clutter"
        #他的名字 如果在类别里面，就复制，如果不在就是clutter
        semantic_gt = np.repeat(class2label[class_name], coords.shape[0])
        semantic_gt = semantic_gt.reshape([-1, 1])
        instance_gt = np.repeat(object_id, coords.shape[0])
        instance_gt = instance_gt.reshape([-1, 1])

        room_coords.append(coords)
        room_colors.append(colors)
        room_semantic_gt.append(semantic_gt)
        room_instance_gt.append(instance_gt)

    room_coords = np.ascontiguousarray(np.vstack(room_coords))
    #
    # if parse_normal:
    #     x_min, z_max, y_min = np.min(room_coords, axis=0)
    #     x_max, z_min, y_max = np.max(room_coords, axis=0)
    #     z_max = -z_max
    #     z_min = -z_min
    #     max_bound = np.array([x_max, y_max, z_max]) + 0.1
    #     min_bound = np.array([x_min, y_min, z_min]) - 0.1
    #     bbox = open3d.geometry.AxisAlignedBoundingBox(
    #         min_bound=min_bound, max_bound=max_bound
    #     )
    #     # crop room
    #     room_mesh = (
    #         area_mesh_dict[os.path.dirname(room)]
    #         .crop(bbox)
    #         .transform(
    #             np.array([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
    #         )
    #     )
    #     vertices = np.array(room_mesh.vertices)
    #     faces = np.array(room_mesh.triangles)
    #     vertex_normals = np.array(room_mesh.vertex_normals)
    #     room_mesh = trimesh.Trimesh(
    #         vertices=vertices, faces=faces, vertex_normals=vertex_normals
    #     )
    #     (closest_points, distances, face_id) = room_mesh.nearest.on_surface(room_coords)
    #     room_normals = room_mesh.face_normals[face_id]

    # if align_angle:
    #     angle = (2 - angle / 180) * np.pi
    #     rot_cos, rot_sin = np.cos(angle), np.sin(angle)
    #     rot_t = np.array([[rot_cos, -rot_sin, 0], [rot_sin, rot_cos, 0], [0, 0, 1]])
    #     room_center = (np.max(room_coords, axis=0) + np.min(room_coords, axis=0)) / 2
    #     room_coords = (room_coords - room_center) @ np.transpose(rot_t) + room_center
    #     if parse_normal:
    #         room_normals = room_normals @ np.transpose(rot_t)

    room_colors = np.ascontiguousarray(np.vstack(room_colors))
    room_semantic_gt = np.ascontiguousarray(np.vstack(room_semantic_gt))
    room_instance_gt = np.ascontiguousarray(np.vstack(room_instance_gt))
    np.save(os.path.join(save_path, "coord.npy"), room_coords.astype(np.float32))
    np.save(os.path.join(save_path, "color.npy"), room_colors.astype(np.uint8))
    np.save(os.path.join(save_path, "segment.npy"), room_semantic_gt.astype(np.int16))
    np.save(os.path.join(save_path, "instance.npy"), room_instance_gt.astype(np.int16))

    if parse_normal:
        np.save(os.path.join(save_path, "normal.npy"), room_normals.astype(np.float32))
class Args:
    #splits = ["Area_1", "Area_2"]  # 指定需要处理的区域
    dataset_root = "/home/qk/PycharmProjects/kpconvX/ml-kpconvx/Standalone/clouds"
    output_root = "/home/qk/PycharmProjects/kpconvX/ml-kpconvx/Standalone/test_output"
    raw_root = None  # 如果不需要可以设为 None
    align_angle = True  # 是否对齐房间角度
    parse_normal = False  # 是否处理法向量
    num_workers = 1  # 预处理时的工作线程数

def main_process():
    # parser = argparse.ArgumentParser()
    # parser.add_argument(
    #     "--splits",
    #     required=True,
    #     nargs="+",
    #     choices=["Area_1", "Area_2", "Area_3", "Area_4", "Area_5", "Area_6"],
    #     help="Splits need to process ([Area_1, Area_2, Area_3, Area_4, Area_5, Area_6]).",
    # )
    # parser.add_argument(
    #     "--dataset_root", required=True, help="Path to Stanford3dDataset_v1.2 dataset"
    # )
    # parser.add_argument(
    #     "--output_root",
    #     required=True,
    #     help="Output path where area folders will be located",
    # )
    # parser.add_argument(
    #     "--raw_root",
    #     default=None,
    #     help="Path to Stanford2d3dDataset_noXYZ dataset (optional)",
    # )
    # parser.add_argument(
    #     "--align_angle", action="store_true", help="Whether align room angles"
    # )
    # parser.add_argument(
    #     "--parse_normal", action="store_true", help="Whether process normal"
    # )
    # parser.add_argument(
    #     "--num_workers", default=1, type=int, help="Num workers for preprocessing."
    # )
    # args = parser.parse_args()



    args = Args()  # 直接创建一个 `args` 对象
    print(args.num_workers)
    # Preprocess data.
    print("Processing scenes...")
    parse_room(args.dataset_root,args.output_root)


if __name__ == "__main__":
    main_process()

其实啥也没改基本，给复制过来了，然后把角度对齐和法向量的都注释掉了，没用上