code/superpoint_ops/blender_visualize_superpoints.py

#!/usr/bin/env python3
"""
S3DIS 超点 / 语义 GT 可视化 —— 仅在 Blender 的 Python（bpy）中运行。

注意：仅含顶点、无面的 PLY 在最终渲染里往往**完全不显示**，PNG 会变成全透明 alpha，
在 IDE 里像灰格「马赛克」。本脚本为每个点生成**面向相机的小三角面片**并写顶点色，
Workbench 才能稳定出图。若需要光滑小球/物理光照，请用 mitsuba_visualize_superpoints.py。

流程：读 coord.npy + 标签 → 在 bpy 中构建带三角面与顶点色的 Mesh
      → 可选保存 .blend → Workbench 渲染 PNG（RGB，无透明底）。

便携 Blender（本仓库已下载）：
  /mnt/data/AODUOLI/PAMI2026/_tools/blender-4.2.2-linux-x64/blender

示例：
  export HOME=/mnt/data/AODUOLI/PAMI2026/_tools/blender_home
  /mnt/data/AODUOLI/PAMI2026/_tools/blender-4.2.2-linux-x64/blender --background \\
    --python /mnt/data/AODUOLI/PAMI2026/blender_visualize_superpoints.py -- \\
    --room Area_1/office_1 \\
    --out_blend /mnt/data/AODUOLI/PAMI2026/outputs/superpoint_vis/office_1_bpy.blend \\
    --out_png /mnt/data/AODUOLI/PAMI2026/outputs/superpoint_vis/office_1_bpy_superpoint.png \\
    --mode superpoint \\
    --max_points 50000

说明：若需要 PNG，请同时传 --out_png；仅看场景文件可只传 --out_blend。
"""
from __future__ import annotations

import argparse
import math
import os
import sys

try:
    import bpy
    from mathutils import Vector
except ImportError:
    print("此脚本必须在 Blender 内运行：blender --background --python ... -- [参数]")
    sys.exit(1)


def _parse_args():
    argv = sys.argv
    if "--" in argv:
        argv = argv[argv.index("--") + 1 :]
    else:
        argv = []
    p = argparse.ArgumentParser()
    root = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "..",
        "_work_biptv3",
        "pointcept_framework",
        "data",
        "s3dis_official",
    )
    p.add_argument("--data_root", type=str, default=os.path.normpath(root))
    p.add_argument("--room", type=str, required=True)
    p.add_argument(
        "--out_blend",
        type=str,
        default=None,
        help="若给定则保存 .blend",
    )
    p.add_argument(
        "--out_png",
        type=str,
        default=None,
        help="若给定则用 Workbench 渲染 PNG（Blender 真·bpy 渲染）",
    )
    p.add_argument(
        "--mode",
        choices=("superpoint", "segment"),
        default="superpoint",
    )
    p.add_argument(
        "--label_npy",
        type=str,
        default=None,
        help="可选：superpoint 模式下覆盖默认的 room/superpoint.npy",
    )
    p.add_argument(
        "--focus_label",
        type=int,
        default=None,
        help="可选：仅高亮某个 superpoint id，其余点置灰，便于查看单个超点形状",
    )
    p.add_argument(
        "--max_points",
        type=int,
        default=None,
        help="随机子采样（建议渲染 PNG 时设 3万～6万）",
    )
    p.add_argument(
        "--temp_ply",
        type=str,
        default=None,
        help="可选：额外写出仅顶点 PLY（调试用；渲染已改为三角网格，不依赖此文件）",
    )
    return p.parse_args(argv)


def _hash_color(uid: int):
    h = (int(uid) * 1103515245 + 12345) & 0x7FFFFFFF
    r = ((h >> 0) & 255) / 255.0
    g = ((h >> 8) & 255) / 255.0
    b = ((h >> 16) & 255) / 255.0
    return r, g, b


S3DIS_CLASSES = [
    "ceiling",
    "floor",
    "wall",
    "beam",
    "column",
    "window",
    "door",
    "table",
    "chair",
    "sofa",
    "bookcase",
    "board",
    "clutter",
]


def _class_color(class_id: int):
    n = len(S3DIS_CLASSES)
    cid = int(class_id) % max(n, 1)
    t = cid / max(n, 1)
    return (
        0.5 + 0.5 * math.cos(2 * math.pi * t),
        0.5 + 0.5 * math.cos(2 * math.pi * t + 2.09),
        0.5 + 0.5 * math.cos(2 * math.pi * t + 4.18),
    )


def _write_ply_binary(path: str, coord, rgb_u8) -> None:
    """coord float32 (N,3), rgb uint8 (N,3) —— Blender ply_import 可识别。"""
    import numpy as np

    n = coord.shape[0]
    header = (
        "ply\nformat binary_little_endian 1.0\n"
        f"element vertex {n}\n"
        "property float x\nproperty float y\nproperty float z\n"
        "property uchar red\nproperty uchar green\nproperty uchar blue\n"
        "end_header\n"
    )
    buf = np.empty(n, dtype=[("x", "<f4"), ("y", "<f4"), ("z", "<f4"), ("r", "u1"), ("g", "u1"), ("b", "u1")])
    buf["x"] = coord[:, 0]
    buf["y"] = coord[:, 1]
    buf["z"] = coord[:, 2]
    buf["r"] = rgb_u8[:, 0]
    buf["g"] = rgb_u8[:, 1]
    buf["b"] = rgb_u8[:, 2]
    with open(path, "wb") as f:
        f.write(header.encode("ascii"))
        f.write(buf.tobytes())


def _camera_from_coord(coord):
    """与原先相机一致；同时返回 billboard 用的 right/up（相机局部 X/Y 轴）。"""
    import numpy as np

    mn = coord.min(axis=0)
    mx = coord.max(axis=0)
    center = (mn + mx) * 0.5
    diag = float(np.linalg.norm(mx - mn)) + 1e-6
    off = Vector((diag * 0.9, -diag * 0.85, diag * 0.55))
    loc = Vector((float(center[0]), float(center[1]), float(center[2]))) + off
    direction = Vector((float(center[0]), float(center[1]), float(center[2]))) - loc
    rot = direction.to_track_quat("-Z", "Y")
    mat = rot.to_matrix().to_4x4()
    mat.translation = loc
    right = mat.col[0].xyz.copy()
    up = mat.col[1].xyz.copy()
    return loc, rot, right, up, diag


def _setup_camera_from_loc_rot(scene, loc: Vector, rot) -> None:
    bpy.ops.object.camera_add(location=loc)
    cam = bpy.context.active_object
    cam.rotation_mode = "QUATERNION"
    cam.rotation_quaternion = rot
    scene.camera = cam
    cam.data.lens = 35


def _build_triangle_point_mesh(coord, rgb_u8, right: Vector, up: Vector, diag: float):
    """每个点三个顶点构成小三角面，顶点色相同；三角位于相机成像平面附近，避免侧面不可见。"""
    import numpy as np

    n = int(coord.shape[0])
    s = float(diag) * 0.0025
    r = np.array([right.x, right.y, right.z], dtype=np.float64)
    u = np.array([up.x, up.y, up.z], dtype=np.float64)
    # 等边三角，中心近似在点上
    v0 = coord + (0.866 * r + 0.0 * u) * s
    v1 = coord + (-0.433 * r + 0.75 * u) * s
    v2 = coord + (-0.433 * r - 0.75 * u) * s
    verts = np.stack([v0, v1, v2], axis=1).reshape(-1, 3)
    faces = [(i * 3, i * 3 + 1, i * 3 + 2) for i in range(n)]

    mesh = bpy.data.meshes.new("S3DIS_points")
    mesh.from_pydata([tuple(v) for v in verts], [], faces)
    mesh.update()
    ca = mesh.color_attributes.new(name="Col", type="FLOAT_COLOR", domain="POINT")
    col = np.repeat(rgb_u8.astype(np.float32) / 255.0, 3, axis=0)
    for i in range(len(ca.data)):
        ca.data[i].color = (float(col[i, 0]), float(col[i, 1]), float(col[i, 2]), 1.0)
    return mesh


def _render_workbench_png(scene, filepath: str) -> None:
    scene.render.engine = "BLENDER_WORKBENCH"
    scene.render.resolution_x = 1920
    scene.render.resolution_y = 1080
    scene.display.shading.light = "STUDIO"
    scene.display.shading.color_type = "VERTEX"
    scene.render.film_transparent = False
    scene.render.image_settings.file_format = "PNG"
    scene.render.image_settings.color_mode = "RGB"
    scene.render.filepath = filepath
    bpy.ops.render.render(write_still=True)


def main():
    import numpy as np

    args = _parse_args()
    if not args.out_blend and not args.out_png:
        sys.exit("请至少指定 --out_blend 或 --out_png 之一")

    room = os.path.join(args.data_root, args.room)
    coord = np.load(os.path.join(room, "coord.npy"))
    n = coord.shape[0]

    if args.mode == "superpoint":
        label_path = args.label_npy or os.path.join(room, "superpoint.npy")
        labels = np.load(label_path).reshape(-1).astype(np.int64)
        color_fn = _hash_color
    else:
        labels = np.load(os.path.join(room, "segment.npy")).reshape(-1).astype(np.int64)
        color_fn = _class_color

    if len(labels) != n:
        raise ValueError("标签长度与 coord 不一致")

    if args.mode == "superpoint":
        print("BLENDER_LABEL_NPY", os.path.abspath(label_path))
        if args.focus_label is not None:
            print("BLENDER_FOCUS_LABEL", int(args.focus_label))

    max_p = args.max_points
    if args.out_png and max_p is None:
        max_p = 50000
    if max_p is not None and n > max_p:
        rng = np.random.default_rng(0)
        if args.mode == "superpoint" and args.focus_label is not None:
            focus_idx = np.flatnonzero(labels == int(args.focus_label))
            if len(focus_idx) == 0:
                raise ValueError("focus_label 不存在于当前标签中")
            if len(focus_idx) >= max_p:
                idx = rng.choice(focus_idx, size=max_p, replace=False)
            else:
                other_idx = np.flatnonzero(labels != int(args.focus_label))
                sample_n = max_p - len(focus_idx)
                sampled_other = rng.choice(other_idx, size=sample_n, replace=False)
                idx = np.concatenate([focus_idx, sampled_other])
        else:
            idx = rng.choice(n, size=max_p, replace=False)
        coord = coord[idx]
        labels = labels[idx]
        n = coord.shape[0]

    rgb = np.zeros((n, 3), dtype=np.float64)
    if args.mode == "superpoint" and args.focus_label is not None:
        rgb[:] = np.array([0.20, 0.20, 0.20], dtype=np.float64)
        focus_mask = labels == int(args.focus_label)
        rgb[focus_mask] = np.array([0.96, 0.34, 0.18], dtype=np.float64)
    else:
        for lid in np.unique(labels):
            c = color_fn(int(lid))
            m = labels == lid
            rgb[m, 0] = c[0]
            rgb[m, 1] = c[1]
            rgb[m, 2] = c[2]
    rgb_u8 = (np.clip(rgb, 0.0, 1.0) * 255.0).astype(np.uint8)
    coord = coord.astype(np.float32)

    here = os.path.dirname(os.path.abspath(__file__))
    out_dir = os.path.join(here, "outputs", "superpoint_vis")
    os.makedirs(out_dir, exist_ok=True)
    if args.temp_ply:
        ply_path = os.path.abspath(args.temp_ply)
        _write_ply_binary(ply_path, coord, rgb_u8)

    loc, rot, right, up, diag = _camera_from_coord(coord)

    bpy.ops.wm.read_factory_settings(use_empty=True)
    mesh = _build_triangle_point_mesh(coord, rgb_u8, right, up, diag)
    obj = bpy.data.objects.new("S3DIS_{}".format(args.mode), mesh)
    bpy.context.collection.objects.link(obj)
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)

    scene = bpy.context.scene
    if args.out_png:
        _setup_camera_from_loc_rot(scene, loc, rot)
        out_png = os.path.abspath(args.out_png)
        os.makedirs(os.path.dirname(out_png) or ".", exist_ok=True)
        _render_workbench_png(scene, out_png)
        print("BLENDER_PNG", out_png)

    if args.out_blend:
        out_b = os.path.abspath(args.out_blend)
        os.makedirs(os.path.dirname(out_b) or ".", exist_ok=True)
        bpy.ops.wm.save_as_mainfile(filepath=out_b)
        print("BLENDER_BLEND", out_b)

    uq = len(np.unique(labels))
    print(
        "[blender] points=",
        n,
        "unique_" + ("superpoints" if args.mode == "superpoint" else "labels"),
        "=",
        uq,
        "mesh=",
        mesh.name,
    )


if __name__ == "__main__":
    main()