code/pointcept_framework/pointcept/datasets/preprocessing/s3dis/preprocess_s3dis.py

"""
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
import pandas as pd
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 generate_superpoint_labels(coords, normals=None, voxel_size=0.12, normal_bins=8):
    """A lightweight geometry-driven superpoint fallback.

    When an external segmentator is unavailable, we build stable region ids by
    combining voxelized coordinates with coarse normal orientation bins.
    """
    coords = np.asarray(coords, dtype=np.float32)
    coord_min = coords.min(axis=0, keepdims=True)
    voxel_coord = np.floor((coords - coord_min) / max(voxel_size, 1e-4)).astype(np.int64)

    if normals is not None and len(normals) == len(coords):
        normals = np.asarray(normals, dtype=np.float32)
        normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-8)
        normal_q = np.floor((normals + 1.0) * 0.5 * normal_bins).astype(np.int64)
        normal_q = np.clip(normal_q, 0, normal_bins)
        tokens = np.concatenate([voxel_coord, normal_q], axis=1)
    else:
        tokens = voxel_coord

    _, inverse = np.unique(tokens, axis=0, return_inverse=True)
    return inverse.astype(np.int32)


def parse_room(
    room, angle, dataset_root, output_root, align_angle=True, parse_normal=False,
    generate_superpoint=False, superpoint_voxel_size=0.12, superpoint_normal_bins=8
):
    print("Parsing: {}".format(room))
    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)
    save_path = os.path.join(output_root, room)
    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 = pd.read_csv(object_path, sep=' ', header=None, usecols=[0,1,2,3,4,5], on_bad_lines='skip', dtype=float, engine='python').to_numpy()
        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"
        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))
    if generate_superpoint:
        room_superpoint = generate_superpoint_labels(
            room_coords,
            room_normals if parse_normal else None,
            voxel_size=superpoint_voxel_size,
            normal_bins=superpoint_normal_bins,
        )
        np.save(os.path.join(save_path, "superpoint.npy"), room_superpoint.astype(np.int32))
    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))


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."
    )
    parser.add_argument(
        "--generate_superpoint", action="store_true",
        help="Whether generate lightweight geometry-aware superpoints."
    )
    parser.add_argument(
        "--superpoint_voxel_size", default=0.12, type=float,
        help="Voxel size used by the fallback superpoint generator."
    )
    parser.add_argument(
        "--superpoint_normal_bins", default=8, type=int,
        help="Number of coarse normal bins used by the fallback superpoint generator."
    )
    args = parser.parse_args()

    if args.parse_normal:
        assert args.raw_root is not None

    room_list = []
    angle_list = []

    # Load room information
    print("Loading room information ...")
    for split in args.splits:
        alignment_file = os.path.join(args.dataset_root, split, f"{split}_alignmentAngle.txt")
        area_info = np.loadtxt(alignment_file, dtype=str)
        if area_info.ndim == 0:
            scalar_value = area_info.item()  # 获取标量字符串
            try:
                # 处理 '[0.0]' 格式：移除括号，转浮点，再整数
                global_angle = int(float(scalar_value.strip('[]')))
            except ValueError:
                global_angle = 0  # 默认 0 如果解析失败
            # 从目录自动列出房间
            room_dirs = sorted([os.path.basename(d) for d in glob.glob(os.path.join(args.dataset_root, split, '*')) if os.path.isdir(d)])
            room_list += [os.path.join(split, room) for room in room_dirs]
            angle_list += [global_angle] * len(room_dirs)
        else:
            room_list += [os.path.join(split, room_info[0]) for room_info in area_info]
            angle_list += [int(room_info[1]) for room_info in area_info]

    if args.parse_normal:
        # load raw mesh file to extract normal
        print("Loading raw mesh file ...")
        for split in args.splits:
            if split != "Area_5":
                mesh_dir = os.path.join(args.raw_root, split, "3d", "rgb.obj")
                mesh = open3d.io.read_triangle_mesh(mesh_dir)
                mesh.triangle_uvs.clear()
            else:
                mesh_a_dir = os.path.join(args.raw_root, f"{split}a", "3d", "rgb.obj")
                mesh_b_dir = os.path.join(args.raw_root, f"{split}b", "3d", "rgb.obj")
                mesh_a = open3d.io.read_triangle_mesh(mesh_a_dir)
                mesh_a.triangle_uvs.clear()
                mesh_b = open3d.io.read_triangle_mesh(mesh_b_dir)
                mesh_b.triangle_uvs.clear()
                mesh_b = mesh_b.transform(
                    np.array(
                        [
                            [0, 0, -1, -4.09703582],
                            [0, 1, 0, 0],
                            [1, 0, 0, -6.22617759],
                            [0, 0, 0, 1],
                        ]
                    )
                )
                mesh = mesh_a + mesh_b
            area_mesh_dict[split] = mesh
            print(f"{split} mesh is loaded")

    # Preprocess data.
    print("Processing scenes...")
    pool = ProcessPoolExecutor(
        max_workers=args.num_workers
    )  # peak 110G memory when parsing normal.
    _ = list(
        pool.map(
            parse_room,
            room_list,
            angle_list,
            repeat(args.dataset_root),
            repeat(args.output_root),
            repeat(args.align_angle),
            repeat(args.parse_normal),
            repeat(args.generate_superpoint),
            repeat(args.superpoint_voxel_size),
            repeat(args.superpoint_normal_bins),
        )
    )


if __name__ == "__main__":
    main_process()