| |
| """ |
| S3DIS 超点可视化（编辑节点 / poplab） |
| |
| 数据侧 superpoint 来源（与预处理一致）： |
| - 每个房间目录下的 superpoint.npy：与 coord.npy 等长的 int32，表示每个点所属超点 id。 |
| - 生成逻辑（无外部分割器时）：体素坐标 +（可选）法向粗分箱 -> np.unique 得到逆映射作为 id。 |
| - 见 pointcept/datasets/preprocessing/s3dis/preprocess_s3dis.py 中 generate_superpoint_labels。 |
| |
| 用法示例： |
| python scripts/visualize_s3dis_superpoints.py \\ |
| --room Area_1/office_1 \\ |
| --out /tmp/superpoints_vis.ply |
| |
| python scripts/visualize_s3dis_superpoints.py --room Area_1/office_1 --recompute --voxel 0.12 --normal-bins 8 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| def generate_superpoint_labels(coords, normals=None, voxel_size=0.12, normal_bins=8): |
| """与 preprocess_s3dis.generate_superpoint_labels 一致的几何 fallback。""" |
| coords = np.asarray(coords, dtype=np.float32) |
| coord_min = coords.min(axis=0, keepdims=True) |
| voxel_coord = np.floor((coords - coord_min) / max(float(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 sp_ids_to_colors(sp: np.ndarray) -> np.ndarray: |
| """稳定伪彩色：对每个 superpoint id 做哈希映射到 [0,1]^3。""" |
| sp = sp.astype(np.int64).reshape(-1) |
| n = len(sp) |
| colors = np.zeros((n, 3), dtype=np.float64) |
| for uid in np.unique(sp): |
| h = (int(uid) * 1103515245 + 12345) & 0x7FFFFFFF |
| r = ((h >> 0) & 255) / 255.0 |
| g = ((h >> 8) & 255) / 255.0 |
| b = ((h >> 16) & 255) / 255.0 |
| colors[sp == uid] = (r, g, b) |
| return np.clip(colors, 0.0, 1.0) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument( |
| "--data_root", |
| type=Path, |
| default=Path(__file__).resolve().parents[1] / "data" / "s3dis_official", |
| help="S3DIS 处理根目录（含 Area_*/*/）", |
| ) |
| ap.add_argument("--room", type=str, required=True, help="例如 Area_1/office_1") |
| ap.add_argument("--out", type=Path, default=Path("superpoints_vis.ply")) |
| ap.add_argument( |
| "--recompute", |
| action="store_true", |
| help="不读 superpoint.npy，按 generate_superpoint_labels 现场重算（用于对照）", |
| ) |
| ap.add_argument("--voxel", type=float, default=0.12) |
| ap.add_argument("--normal-bins", type=int, default=8) |
| args = ap.parse_args() |
|
|
| room_dir = args.data_root / args.room |
| coord_p = room_dir / "coord.npy" |
| if not coord_p.is_file(): |
| raise FileNotFoundError(coord_p) |
|
|
| coord = np.load(coord_p) |
| normal = None |
| npy_n = room_dir / "normal.npy" |
| if npy_n.is_file(): |
| normal = np.load(npy_n) |
|
|
| if args.recompute: |
| sp = generate_superpoint_labels( |
| coord, normals=normal, voxel_size=args.voxel, normal_bins=args.normal_bins |
| ) |
| tag = "recomputed" |
| else: |
| sp_p = room_dir / "superpoint.npy" |
| if not sp_p.is_file(): |
| raise FileNotFoundError( |
| f"{sp_p} 不存在，可先跑预处理加 --generate_superpoint，或改用 --recompute" |
| ) |
| sp = np.load(sp_p).reshape(-1).astype(np.int32) |
| tag = "disk" |
|
|
| if len(sp) != len(coord): |
| raise ValueError(f"superpoint 长度 {len(sp)} != coord 长度 {len(coord)}") |
|
|
| colors = sp_ids_to_colors(sp) |
| try: |
| import open3d as o3d |
|
|
| pcd = o3d.geometry.PointCloud() |
| pcd.points = o3d.utility.Vector3dVector(coord.astype(np.float64)) |
| pcd.colors = o3d.utility.Vector3dVector(colors) |
| args.out.parent.mkdir(parents=True, exist_ok=True) |
| o3d.io.write_point_cloud(str(args.out), pcd) |
| print(f"[{tag}] wrote {args.out} points={len(coord)} unique_superpoints={len(np.unique(sp))}") |
| print("若在本机有显示器，可取消注释 o3d.visualization.draw_geometries([pcd])") |
| |
| except ImportError: |
| print("未安装 open3d，无法写 PLY；请: pip install open3d") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
