code/superpoint_ops/superpoint_delivery_bundle.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import colorsys
import json
import os
import subprocess
import tempfile
from pathlib import Path

import numpy as np
from scipy.spatial import cKDTree


_RENDER_HELPER = r"""
from __future__ import annotations
import argparse
import json
import os
import sys
try:
    import bpy
    from mathutils import Vector
except ImportError:
    print('This script must run inside Blender.')
    sys.exit(1)

def parse_args():
    argv = sys.argv
    argv = argv[argv.index('--') + 1:] if '--' in argv else []
    ap = argparse.ArgumentParser(description='Render NPZ pointcloud assets with Blender Workbench')
    ap.add_argument('--asset_npz', type=str, required=True)
    ap.add_argument('--out_png', type=str, required=True)
    ap.add_argument('--out_blend', type=str, default=None)
    ap.add_argument('--background', choices=('white', 'light', 'dark'), default=None)
    return ap.parse_args(argv)

def camera_from_coord(coord):
    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 build_triangle_mesh(coord, rgb_u8, right, up, diag):
    import numpy as np, math
    n = int(coord.shape[0])
    scale = float(diag) * 0.0042
    r = np.array([right.x, right.y, right.z], dtype=np.float64)
    u = np.array([up.x, up.y, up.z], dtype=np.float64)
    # 用 12 边形替代六边形，减轻“尖尖的”片元边缘
    poly_n = 12
    poly_offsets = []
    for k in range(poly_n):
        angle = 2.0 * math.pi * k / poly_n
        poly_offsets.append(math.cos(angle) * r + math.sin(angle) * u)
    verts_per_pt = poly_n + 1  # center + polygon corners
    all_verts = np.empty((n * verts_per_pt, 3), dtype=np.float64)
    all_faces = []
    for k in range(poly_n):
        off = poly_offsets[k] * scale
        all_verts[np.arange(n) * verts_per_pt + (k + 1)] = coord + off
    all_verts[np.arange(n) * verts_per_pt] = coord  # center vertices
    for i in range(n):
        base = i * verts_per_pt
        for k in range(poly_n):
            c1 = base + 1 + k
            c2 = base + 1 + (k + 1) % poly_n
            all_faces.append((base, c1, c2))
    mesh = bpy.data.meshes.new('StoryPoints')
    mesh.from_pydata([tuple(v) for v in all_verts], [], all_faces)
    mesh.update()
    attr = mesh.color_attributes.new(name='Col', type='FLOAT_COLOR', domain='POINT')
    col = (rgb_u8.astype(np.float32) / 255.0).repeat(verts_per_pt, axis=0)
    for i in range(len(attr.data)):
        attr.data[i].color = (float(col[i, 0]), float(col[i, 1]), float(col[i, 2]), 1.0)
    return mesh

def setup_camera(scene, loc, rot):
    bpy.ops.object.camera_add(location=loc)
    cam = bpy.context.active_object
    cam.rotation_mode = 'QUATERNION'
    cam.rotation_quaternion = rot
    cam.data.lens = 35
    scene.camera = cam

def set_background(scene, mode):
    world = scene.world or bpy.data.worlds.new('World')
    scene.world = world
    world.use_nodes = True
    bg = world.node_tree.nodes.get('Background')
    if bg is None:
        bg = world.node_tree.nodes.new(type='ShaderNodeBackground')
    if mode == 'dark':
        bg.inputs[0].default_value = (0.02, 0.02, 0.025, 1.0)
    elif mode == 'light':
        bg.inputs[0].default_value = (0.97, 0.97, 0.975, 1.0)
    else:
        bg.inputs[0].default_value = (1.0, 1.0, 1.0, 1.0)
    bg.inputs[1].default_value = 1.0

def render_png(scene, path, background):
    scene.render.engine = 'BLENDER_WORKBENCH'
    # Increase native render resolution to soften jagged point boundaries.
    scene.render.resolution_x = 2560
    scene.render.resolution_y = 1440
    scene.render.image_settings.file_format = 'PNG'
    scene.render.image_settings.color_mode = 'RGBA'
    scene.render.film_transparent = True
    scene.view_settings.view_transform = 'Standard'
    scene.view_settings.look = 'None'
    scene.display.shading.color_type = 'VERTEX'
    scene.display.shading.light = 'FLAT'
    scene.display.shading.show_cavity = False
    scene.display.shading.show_object_outline = False
    scene.display.shading.show_specular_highlight = False
    scene.render.filepath = path
    bpy.ops.render.render(write_still=True)

def main():
    import numpy as np
    args = parse_args()
    asset = np.load(args.asset_npz)
    coord = asset['coord'].astype(np.float32)
    rgb = asset['rgb'].astype(np.uint8)
    meta_path = os.path.splitext(args.asset_npz)[0] + '.json'
    meta = {}
    if os.path.isfile(meta_path):
        meta = json.loads(open(meta_path, 'r', encoding='utf-8').read())
    bg = args.background or meta.get('background', 'white')
    bpy.ops.wm.read_factory_settings(use_empty=True)
    loc, rot, right, up, diag = camera_from_coord(coord)
    mesh = build_triangle_mesh(coord, rgb, right, up, diag)
    obj = bpy.data.objects.new('StoryObject', mesh)
    bpy.context.collection.objects.link(obj)
    setup_camera(bpy.context.scene, loc, rot)
    set_background(bpy.context.scene, bg)
    out_png = os.path.abspath(args.out_png)
    os.makedirs(os.path.dirname(out_png) or '.', exist_ok=True)
    render_png(bpy.context.scene, out_png, bg)
    print('BLENDER_PNG', out_png)
    if args.out_blend:
        out_blend = os.path.abspath(args.out_blend)
        os.makedirs(os.path.dirname(out_blend) or '.', exist_ok=True)
        bpy.ops.wm.save_as_mainfile(filepath=out_blend)
        print('BLENDER_BLEND', out_blend)

if __name__ == '__main__':
    main()
"""

DEFAULT_DATA_ROOT = Path('/mnt/data/AODUOLI/_work_biptv3/pointcept_framework/data/s3dis_official')
DEFAULT_BLENDER_BIN = Path('/mnt/data/AODUOLI/PAMI2026/_tools/blender-4.2.2-linux-x64/blender')
DEFAULT_SEGMENTATOR_CPP = Path('/mnt/data/AODUOLI/PAMI2026/_codex_research/segmentator/csrc/segmentator.cpp')
DEFAULT_SEGMENTATOR_BUILD_DIR = Path('/mnt/data/AODUOLI/PAMI2026/_codex_research/segmentator/csrc/torch_build')


def stable_hash_color(uid: int) -> np.ndarray:
    import colorsys
    h = (int(uid) * 1103515245 + 12345) & 0x7FFFFFFF
    hue = (h & 0xFFFF) / 65536.0
    sat = 0.75 + ((h >> 16) & 0xFF) / 255.0 * 0.25  # 0.75-1.00
    val = 0.40 + ((h >> 24) & 0x7F) / 127.0 * 0.30  # 0.40-0.70
    r, g, b = colorsys.hsv_to_rgb(hue, sat, val)
    return np.array([int(r * 255), int(g * 255), int(b * 255)], dtype=np.uint8)


S3DIS_SEGMENT_RGB_U8 = np.array([
    [230, 230, 235],  # ceiling
    [184, 133, 77],   # floor
    [107, 184, 122],  # wall
    [89, 140, 217],   # beam
    [242, 191, 64],   # column
    [115, 217, 242],  # window
    [235, 89, 89],    # door
    [242, 140, 51],   # table
    [191, 115, 235],  # chair
    [140, 102, 224],  # sofa
    [128, 209, 115],  # bookcase
    [224, 128, 166],  # board
    [140, 140, 148],  # clutter
], dtype=np.uint8)


def _clip_rgb01(rgb) -> tuple[float, float, float]:
    return tuple(float(np.clip(c, 0.0, 1.0)) for c in rgb)


def semantic_family_color(segment_id: int, superpoint_id: int, theme: str) -> np.ndarray:
    """Same semantic label -> same fixed semantic palette, with mild intra-class variation."""
    sid = int(segment_id)
    if 0 <= sid < len(S3DIS_SEGMENT_RGB_U8):
        base_rgb = S3DIS_SEGMENT_RGB_U8[sid].astype(np.float32) / 255.0
    else:
        hue = ((sid * 0.137) % 1.0)
        base_rgb = np.array(colorsys.hsv_to_rgb(hue, 0.45, 0.86), dtype=np.float32)
    h, s, v = colorsys.rgb_to_hsv(*_clip_rgb01(base_rgb))
    jitter = (int(superpoint_id) * 1103515245 + 12345) & 0x7FFFFFFF
    if theme == 'dark':
        sat = np.clip(s * (0.94 + ((jitter >> 8) & 0xFF) / 255.0 * 0.16), 0.28, 0.88)
        val = np.clip(v * (0.92 + ((jitter >> 16) & 0xFF) / 255.0 * 0.16), 0.55, 0.96)
    elif theme == 'light':
        sat = np.clip(s * (0.90 + ((jitter >> 8) & 0xFF) / 255.0 * 0.12), 0.18, 0.80)
        val = np.clip(v * (0.95 + ((jitter >> 16) & 0xFF) / 255.0 * 0.10), 0.68, 0.98)
    else:
        sat = np.clip(s * (0.92 + ((jitter >> 8) & 0xFF) / 255.0 * 0.12), 0.22, 0.84)
        val = np.clip(v * (0.96 + ((jitter >> 16) & 0xFF) / 255.0 * 0.08), 0.74, 0.98)
    rgb = colorsys.hsv_to_rgb(h, sat, val)
    return np.array([int(c * 255) for c in _clip_rgb01(rgb)], dtype=np.uint8)


def canonical_generate_superpoints(coords: np.ndarray, normals: np.ndarray | None = None, voxel_size: float = 0.12, normal_bins: int = 8) -> np.ndarray:
    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 load_segmentator_module(segmentator_cpp: Path, build_dir: Path, verbose: bool = False):
    from torch.utils.cpp_extension import load

    segmentator_cpp = Path(segmentator_cpp)
    build_dir = Path(build_dir)
    if not segmentator_cpp.is_file():
        raise FileNotFoundError(f'Segmentator source not found: {segmentator_cpp}')
    build_dir.mkdir(parents=True, exist_ok=True)
    return load(
        name='libsegmentator_dyn',
        sources=[str(segmentator_cpp)],
        build_directory=str(build_dir),
        extra_cflags=['-O3'],
        verbose=verbose,
    )


def segmentator_generate_superpoints(
    coords: np.ndarray,
    normals: np.ndarray,
    knn_k: int = 50,
    k_thresh: float = 0.01,
    seg_min_verts: int = 20,
    segmentator_cpp: Path = DEFAULT_SEGMENTATOR_CPP,
    build_dir: Path = DEFAULT_SEGMENTATOR_BUILD_DIR,
    build_verbose: bool = False,
) -> np.ndarray:
    import torch
    from torch_cluster import knn_graph

    coords = np.ascontiguousarray(np.asarray(coords, dtype=np.float32))
    normals = np.ascontiguousarray(np.asarray(normals, dtype=np.float32))
    if coords.shape != normals.shape:
        raise ValueError(f'coords/normals mismatch: {coords.shape} vs {normals.shape}')
    normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-8)

    seg = load_segmentator_module(segmentator_cpp=segmentator_cpp, build_dir=build_dir, verbose=build_verbose)
    pts = torch.from_numpy(coords)
    nrm = torch.from_numpy(normals)
    edges = knn_graph(pts, k=int(knn_k)).T.contiguous().to(dtype=torch.int64, device='cpu')
    labels = seg.segment_point(pts.contiguous(), nrm.contiguous(), edges, float(k_thresh), int(seg_min_verts))
    labels = labels.cpu().numpy().reshape(-1).astype(np.int64)
    _, inverse = np.unique(labels, return_inverse=True)
    return inverse.astype(np.int32)


def sample_indices(n: int, max_points: int | None, rng: np.random.Generator) -> np.ndarray:
    if max_points is None or n <= max_points:
        return np.arange(n, dtype=np.int64)
    idx = rng.choice(n, size=max_points, replace=False)
    return np.sort(idx.astype(np.int64))


def knn_boundary_mask(coords: np.ndarray, labels: np.ndarray, k: int, diff_ratio: float) -> np.ndarray:
    if len(coords) <= 3:
        return np.zeros(len(coords), dtype=bool)
    k_eff = min(int(k) + 1, len(coords))
    tree = cKDTree(coords)
    _, nn = tree.query(coords, k=k_eff)
    if nn.ndim == 1:
        nn = nn[:, None]
    nbr = nn[:, 1:]
    if nbr.size == 0:
        return np.zeros(len(coords), dtype=bool)
    diff = labels[nbr] != labels[:, None]
    return diff.mean(axis=1) >= float(diff_ratio)


def knn_shell_mask(coords: np.ndarray, labels: np.ndarray, focus_label: int, k: int = 26) -> np.ndarray:
    focus = labels == int(focus_label)
    if not np.any(focus):
        return np.zeros(len(coords), dtype=bool)
    k_eff = min(int(k) + 1, len(coords))
    tree = cKDTree(coords)
    _, nn = tree.query(coords[focus], k=k_eff)
    if nn.ndim == 1:
        nn = nn[:, None]
    shell = np.zeros(len(coords), dtype=bool)
    shell[np.unique(nn[:, 1:].reshape(-1))] = True
    shell[focus] = False
    return shell


def auto_focus_label(coords: np.ndarray, labels: np.ndarray) -> int:
    uniq, inverse, counts = np.unique(labels, return_inverse=True, return_counts=True)
    if uniq.size == 1:
        return int(uniq[0])
    centroids = np.zeros((uniq.size, 3), dtype=np.float64)
    np.add.at(centroids, inverse, coords)
    centroids /= counts[:, None]
    scene_center = coords.mean(axis=0, keepdims=True)
    lo = max(120, np.percentile(counts, 60))
    hi = max(lo, np.percentile(counts, 95))
    candidate = (counts >= lo) & (counts <= hi)
    if not candidate.any():
        candidate = counts >= np.median(counts)
    dist = np.linalg.norm(centroids - scene_center, axis=1)
    score = dist / np.log1p(counts)
    score = np.where(candidate, score, np.inf)
    return int(uniq[int(np.argmin(score))])


def local_focus_indices(coords: np.ndarray, focus_mask: np.ndarray, limit: int) -> np.ndarray:
    if not np.any(focus_mask):
        return np.arange(min(limit, len(coords)), dtype=np.int64)
    center = coords[focus_mask].mean(axis=0)
    tree = cKDTree(coords)
    _, idx = tree.query(center, k=min(int(limit), len(coords)))
    idx = np.atleast_1d(idx).astype(np.int64)
    focus_idx = np.flatnonzero(focus_mask)
    merged = np.unique(np.concatenate([idx, focus_idx]))
    if len(merged) > limit:
        d = np.linalg.norm(coords[merged] - center[None, :], axis=1)
        merged = merged[np.argsort(d)[:limit]]
    return np.sort(merged.astype(np.int64))


def cutaway_indices(
    coords: np.ndarray,
    max_points: int,
    rng: np.random.Generator,
    ceiling_ratio: float = 0.86,
    wall_depth: float = 0.12,
    walls: tuple[str, ...] = ('xmin', 'ymin', 'xmax'),
) -> np.ndarray:
    n = len(coords)
    if n == 0:
        return np.empty((0,), dtype=np.int64)
    pmin = coords.min(axis=0)
    pmax = coords.max(axis=0)
    prange = np.maximum(pmax - pmin, 1e-6)
    keep = coords[:, 2] < (pmin[2] + prange[2] * float(ceiling_ratio))
    for wall in walls:
        if wall == 'xmin':
            keep &= coords[:, 0] > (pmin[0] + prange[0] * float(wall_depth))
        elif wall == 'xmax':
            keep &= coords[:, 0] < (pmax[0] - prange[0] * float(wall_depth))
        elif wall == 'ymin':
            keep &= coords[:, 1] > (pmin[1] + prange[1] * float(wall_depth))
        elif wall == 'ymax':
            keep &= coords[:, 1] < (pmax[1] - prange[1] * float(wall_depth))
    idx = np.flatnonzero(keep)
    if len(idx) == 0:
        return sample_indices(n, max_points, rng)
    if len(idx) > max_points:
        idx = rng.choice(idx, size=max_points, replace=False)
    return np.sort(idx.astype(np.int64))


def build_rgb(
    mode: str,
    superpoint: np.ndarray,
    segment: np.ndarray,
    sp_boundary: np.ndarray,
    seg_boundary: np.ndarray,
    focus_label: int,
    focus_shell: np.ndarray,
    theme: str,
) -> np.ndarray:
    rgb = np.zeros((len(superpoint), 3), dtype=np.uint8)
    dark = theme == 'dark'

    if mode == 'superpoint':
        for u in np.unique(superpoint):
            rgb[superpoint == u] = stable_hash_color(int(u))
        return rgb
    if mode == 'semantic_family':
        for u in np.unique(superpoint):
            mask = superpoint == u
            seg_vals = segment[mask].astype(np.int64)
            seg_vals = seg_vals[seg_vals >= 0]
            dom_seg = int(np.bincount(seg_vals).argmax()) if seg_vals.size else -1
            rgb[mask] = semantic_family_color(dom_seg, int(u), theme)
        return rgb

    if dark:
        base = np.array([40, 40, 45], dtype=np.uint8)
        neutral = np.array([28, 28, 34], dtype=np.uint8)
        focus_bg = np.array([20, 20, 24], dtype=np.uint8)
        same_seg = np.array([66, 82, 95], dtype=np.uint8)
    elif theme == 'light':
        base = np.array([235, 235, 238], dtype=np.uint8)
        neutral = np.array([245, 245, 247], dtype=np.uint8)
        focus_bg = np.array([248, 248, 250], dtype=np.uint8)
        same_seg = np.array([218, 223, 229], dtype=np.uint8)
    else:
        base = np.array([242, 242, 242], dtype=np.uint8)
        neutral = np.array([250, 250, 250], dtype=np.uint8)
        focus_bg = np.array([252, 252, 252], dtype=np.uint8)
        same_seg = np.array([224, 229, 235], dtype=np.uint8)

    orange = np.array([236, 141, 53], dtype=np.uint8)
    cyan = np.array([56, 168, 240], dtype=np.uint8)
    red = np.array([220, 72, 60], dtype=np.uint8)
    shell = np.array([95, 197, 226], dtype=np.uint8)
    focus = np.array([236, 96, 67], dtype=np.uint8)

    if mode == 'sp_boundary':
        rgb[:] = base
        rgb[sp_boundary] = orange
        return rgb
    if mode == 'segment_boundary':
        rgb[:] = base
        rgb[seg_boundary] = cyan
        return rgb
    if mode == 'boundary_relation':
        both = sp_boundary & seg_boundary
        sp_only = sp_boundary & ~seg_boundary
        seg_only = ~sp_boundary & seg_boundary
        rgb[:] = neutral
        rgb[seg_only] = cyan
        rgb[sp_only] = orange
        rgb[both] = red
        return rgb
    if mode == 'focus_superpoint':
        focus_mask = superpoint == int(focus_label)
        same_seg_mask = np.zeros_like(focus_mask)
        if np.any(focus_mask):
            dom_seg = int(np.bincount(segment[focus_mask]).argmax())
            same_seg_mask = segment == dom_seg
        rgb[:] = focus_bg
        rgb[same_seg_mask] = same_seg
        rgb[focus_shell] = shell
        rgb[focus_mask] = focus
        return rgb
    raise ValueError(mode)


def save_asset(path: Path, coord: np.ndarray, rgb: np.ndarray, meta: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(path, coord=coord.astype(np.float32), rgb=rgb.astype(np.uint8))
    path.with_suffix('.json').write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding='utf-8')


def build_story_grid(room_out: Path, room_tag: str, theme: str) -> Path | None:
    try:
        from PIL import Image, ImageDraw
    except Exception:
        return None
    order = [
        ('semantic_family', 'Semantic-Family Superpoints'),
        ('semantic_family_cutaway', 'Semantic-Family Cutaway'),
        ('superpoint', 'Raw-ID Superpoints'),
        ('superpoint_cutaway', 'Raw-ID Cutaway'),
        ('sp_boundary', 'Superpoint Boundary'),
        ('segment_boundary', 'Semantic Boundary'),
        ('boundary_relation', 'Boundary Relation'),
        ('focus_superpoint', 'Focus Superpoint'),
    ]
    imgs = []
    for suffix, title in order:
        p = room_out / f'{room_tag}_{suffix}.png'
        if p.is_file():
            imgs.append((title, Image.open(p).convert('RGB')))
    if not imgs:
        return None
    card_w = 960
    card_h = 600
    pad = 24
    cols = 2
    rows = (len(imgs) + cols - 1) // cols
    if theme == 'dark':
        canvas_bg = (18, 18, 22)
        panel_bg = (28, 28, 34)
        title_fg = (235, 235, 240)
    elif theme == 'light':
        canvas_bg = (245, 245, 247)
        panel_bg = (252, 252, 253)
        title_fg = (38, 38, 45)
    else:
        canvas_bg = (255, 255, 255)
        panel_bg = (250, 250, 250)
        title_fg = (25, 25, 28)
    canvas = Image.new('RGB', (cols * card_w + (cols + 1) * pad, rows * card_h + (rows + 1) * pad), canvas_bg)
    for idx, (title, im) in enumerate(imgs):
        r = idx // cols
        c = idx % cols
        x = pad + c * (card_w + pad)
        y = pad + r * (card_h + pad)
        panel = Image.new('RGB', (card_w, card_h), panel_bg)
        thumb = im.copy()
        thumb.thumbnail((card_w, card_h - 36))
        ox = (card_w - thumb.width) // 2
        oy = 36 + (card_h - 36 - thumb.height) // 2
        panel.paste(thumb, (ox, oy))
        d = ImageDraw.Draw(panel)
        d.text((18, 10), title, fill=title_fg)
        canvas.paste(panel, (x, y))
    out = room_out / f'{room_tag}_story_grid.png'
    canvas.save(out)
    return out


def resolve_room_dir(data_root: Path, room: str) -> Path:
    room_dir = data_root / room
    if not room_dir.is_dir():
        raise FileNotFoundError(f'Room not found: {room_dir}')
    return room_dir


def load_superpoint_labels(
    room_dir: Path,
    source: str,
    label_npy: Path | None,
    voxel_size: float,
    normal_bins: int,
    knn_k: int,
    k_thresh: float,
    seg_min_verts: int,
    segmentator_cpp: Path,
    segmentator_build_dir: Path,
    build_verbose: bool,
) -> tuple[np.ndarray, str]:
    coord = np.load(room_dir / 'coord.npy').astype(np.float32)
    normal_file = room_dir / 'normal.npy'
    normals = np.load(normal_file).astype(np.float32) if normal_file.is_file() else None

    if label_npy is not None:
        labels = np.load(label_npy).reshape(-1).astype(np.int64)
        return labels, str(label_npy)
    if source == 'existing':
        p = room_dir / 'superpoint.npy'
        labels = np.load(p).reshape(-1).astype(np.int64)
        return labels, str(p)
    if source == 'canonical':
        labels = canonical_generate_superpoints(coord, normals, voxel_size=voxel_size, normal_bins=normal_bins).astype(np.int64)
        return labels, f'canonical(voxel={voxel_size}, normal_bins={normal_bins})'
    if source == 'segmentator':
        if normals is None:
            raise FileNotFoundError(f'normal.npy is required for segmentator source: {room_dir}')
        labels = segmentator_generate_superpoints(
            coord,
            normals,
            knn_k=knn_k,
            k_thresh=k_thresh,
            seg_min_verts=seg_min_verts,
            segmentator_cpp=segmentator_cpp,
            build_dir=segmentator_build_dir,
            build_verbose=build_verbose,
        ).astype(np.int64)
        return labels, f'segmentator(knn_k={knn_k}, k_thresh={k_thresh}, seg_min_verts={seg_min_verts})'
    raise ValueError(f'Unknown source: {source}')


def write_generated_labels(room_dir: Path, labels: np.ndarray, output_root: Path) -> Path:
    out = output_root / room_dir.parent.name / room_dir.name / 'superpoint.npy'
    out.parent.mkdir(parents=True, exist_ok=True)
    np.save(str(out), labels.astype(np.int32))
    return out


def prepare_story_assets(
    coord: np.ndarray,
    segment: np.ndarray,
    superpoint: np.ndarray,
    label_desc: str,
    room: str,
    out_dir: Path,
    max_points: int,
    boundary_knn: int,
    sp_boundary_ratio: float,
    seg_boundary_ratio: float,
    focus_knn: int,
    focus_label: int | None,
    focus_view_points: int,
    modes: list[str],
    background_theme: str,
) -> dict:
    rng = np.random.default_rng(0)
    base_idx = sample_indices(len(coord), max_points, rng)
    cut_idx = cutaway_indices(coord, max_points, rng)
    coord_b = coord[base_idx]
    segment_b = segment[base_idx]
    superpoint_b = superpoint[base_idx]
    coord_c = coord[cut_idx]
    segment_c = segment[cut_idx]
    superpoint_c = superpoint[cut_idx]
    sp_boundary_b = knn_boundary_mask(coord_b, superpoint_b, k=boundary_knn, diff_ratio=sp_boundary_ratio)
    seg_boundary_b = knn_boundary_mask(coord_b, segment_b, k=boundary_knn, diff_ratio=seg_boundary_ratio)
    sp_boundary_c = knn_boundary_mask(coord_c, superpoint_c, k=boundary_knn, diff_ratio=sp_boundary_ratio)
    seg_boundary_c = knn_boundary_mask(coord_c, segment_c, k=boundary_knn, diff_ratio=seg_boundary_ratio)
    overlap_b = sp_boundary_b & seg_boundary_b
    overlap_c = sp_boundary_c & seg_boundary_c
    focus_label = int(focus_label) if focus_label is not None else auto_focus_label(coord, superpoint)
    room_tag = room.replace('/', '_')
    summary = {
        'room': room,
        'label_source': label_desc,
        'background_theme': background_theme,
        'points_full': int(len(coord)),
        'points_sampled': int(len(coord_b)),
        'points_cutaway_sampled': int(len(coord_c)),
        'unique_superpoints_full': int(np.unique(superpoint).size),
        'unique_superpoints_sampled': int(np.unique(superpoint_b).size),
        'unique_superpoints_cutaway_sampled': int(np.unique(superpoint_c).size),
        'sp_boundary_points_sampled': int(sp_boundary_b.sum()),
        'segment_boundary_points_sampled': int(seg_boundary_b.sum()),
        'boundary_overlap_points_sampled': int(overlap_b.sum()),
        'sp_boundary_points_cutaway_sampled': int(sp_boundary_c.sum()),
        'segment_boundary_points_cutaway_sampled': int(seg_boundary_c.sum()),
        'boundary_overlap_points_cutaway_sampled': int(overlap_c.sum()),
        'cutaway': {
            'ceiling_ratio': 0.86,
            'wall_depth': 0.12,
            'walls_removed': ['xmin', 'ymin', 'xmax'],
        },
        'focus_label': int(focus_label),
        'focus_label_points_full': int((superpoint == focus_label).sum()),
        'boundary_knn': int(boundary_knn),
        'sp_boundary_ratio': float(sp_boundary_ratio),
        'seg_boundary_ratio': float(seg_boundary_ratio),
        'boundary_definition': 'A point is marked as a boundary point when at least diff_ratio of its k nearest neighbors belong to a different label.',
        'modes': list(modes),
    }
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / f'{room_tag}_summary.json').write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding='utf-8')
    assets_dir = out_dir / 'assets'
    assets_dir.mkdir(parents=True, exist_ok=True)

    for mode in modes:
        base_mode = mode[:-8] if mode.endswith('_cutaway') else mode
        use_cutaway = mode.endswith('_cutaway')
        if mode == 'focus_superpoint':
            focus_mask_full = superpoint == focus_label
            focus_idx = local_focus_indices(coord, focus_mask_full, focus_view_points)
            coord_m = coord[focus_idx]
            segment_m = segment[focus_idx]
            superpoint_m = superpoint[focus_idx]
            focus_shell = knn_shell_mask(coord_m, superpoint_m, focus_label, k=focus_knn)
            rgb = build_rgb(mode, superpoint_m, segment_m, np.zeros(len(coord_m), dtype=bool), np.zeros(len(coord_m), dtype=bool), focus_label, focus_shell, background_theme)
            meta = dict(summary)
            meta.update({'mode': mode, 'background': background_theme, 'sampled_points': int(len(coord_m))})
            save_asset(assets_dir / f'{room_tag}_{mode}.npz', coord_m, rgb, meta)
            continue
        if use_cutaway:
            coord_m = coord_c
            segment_m = segment_c
            superpoint_m = superpoint_c
            sp_boundary_m = sp_boundary_c
            seg_boundary_m = seg_boundary_c
        else:
            coord_m = coord_b
            segment_m = segment_b
            superpoint_m = superpoint_b
            sp_boundary_m = sp_boundary_b
            seg_boundary_m = seg_boundary_b
        rgb = build_rgb(base_mode, superpoint_m, segment_m, sp_boundary_m, seg_boundary_m, focus_label, np.zeros(len(coord_m), dtype=bool), background_theme)
        meta = dict(summary)
        meta.update({'mode': mode, 'background': background_theme, 'sampled_points': int(len(coord_m))})
        save_asset(assets_dir / f'{room_tag}_{mode}.npz', coord_m, rgb, meta)
    return summary


def _composite_white_bg(png_path: Path, bg_color: tuple = (255, 255, 255)) -> None:
    """将 RGBA PNG 合成到纯色背景上，覆盖原文件。"""
    try:
        from PIL import Image
        im = Image.open(png_path)
        if im.mode != 'RGBA':
            return
        bg = Image.new('RGB', im.size, bg_color)
        bg.paste(im, mask=im.split()[3])
        bg.save(png_path)
    except Exception as e:
        print(f'composite failed for {png_path}: {e}')


def render_assets(room_out: Path, blender_bin: Path, background_theme: str = 'white') -> list[Path]:
    assets = sorted((room_out / 'assets').glob('*.npz'))
    if not assets:
        return []
    if not blender_bin.exists():
        raise FileNotFoundError(f'Blender not found: {blender_bin}')
    os.environ.setdefault('HOME', str(Path(room_out).parents[2] / '_tools' / 'blender_home'))
    Path(os.environ['HOME']).mkdir(parents=True, exist_ok=True)
    helper_file = Path(tempfile.mkdtemp(prefix='sp_bundle_')) / 'render_helper.py'
    helper_file.write_text(_RENDER_HELPER, encoding='utf-8')
    outputs = []
    try:
        for asset in assets:
            name = asset.stem
            out_png = room_out / f'{name}.png'
            cmd = [str(blender_bin), '--background', '--python', str(helper_file), '--', '--asset_npz', str(asset), '--out_png', str(out_png)]
            subprocess.run(cmd, check=True)
            if background_theme != 'dark':
                bg = (255, 255, 255) if background_theme == 'white' else (247, 247, 249)
                _composite_white_bg(out_png, bg)
            outputs.append(out_png)
    finally:
        try:
            helper_file.unlink(missing_ok=True)
            helper_file.parent.rmdir()
        except Exception:
            pass
    return outputs


def cmd_generate(args: argparse.Namespace) -> None:
    room_dir = resolve_room_dir(args.data_root, args.room)
    coord = np.load(room_dir / 'coord.npy').astype(np.float32)
    normals = None
    normal_file = room_dir / 'normal.npy'
    if normal_file.is_file():
        normals = np.load(normal_file).astype(np.float32)

    if args.method == 'canonical':
        labels = canonical_generate_superpoints(coord, normals, voxel_size=args.voxel_size, normal_bins=args.normal_bins)
    else:
        if normals is None:
            raise FileNotFoundError(f'normal.npy is required for segmentator method: {room_dir}')
        labels = segmentator_generate_superpoints(
            coord,
            normals,
            knn_k=args.knn_k,
            k_thresh=args.k_thresh,
            seg_min_verts=args.seg_min_verts,
            segmentator_cpp=args.segmentator_cpp,
            build_dir=args.segmentator_build_dir,
            build_verbose=args.build_verbose,
        )
    out = write_generated_labels(room_dir, labels, args.output_root)
    uniq, cnt = np.unique(labels, return_counts=True)
    print(json.dumps({
        'room': args.room,
        'method': args.method,
        'output': str(out),
        'num_superpoints': int(uniq.size),
        'mean_points_per_superpoint': float(cnt.mean()),
        'max_points_per_superpoint': int(cnt.max()),
    }, ensure_ascii=False))


def cmd_story(args: argparse.Namespace) -> None:
    for room in args.rooms:
        room_dir = resolve_room_dir(args.data_root, room)
        coord = np.load(room_dir / 'coord.npy').astype(np.float32)
        segment = np.load(room_dir / 'segment.npy').reshape(-1).astype(np.int64)
        superpoint, label_desc = load_superpoint_labels(
            room_dir,
            args.source,
            args.label_npy,
            args.voxel_size,
            args.normal_bins,
            args.knn_k,
            args.k_thresh,
            args.seg_min_verts,
            args.segmentator_cpp,
            args.segmentator_build_dir,
            args.build_verbose,
        )
        room_tag = room.replace('/', '_')
        room_out = args.output_root / room_tag
        summary = prepare_story_assets(
            coord,
            segment,
            superpoint,
            label_desc,
            room,
            room_out,
            args.max_points,
            args.boundary_knn,
            args.sp_boundary_ratio,
            args.seg_boundary_ratio,
            args.focus_knn,
            args.focus_label,
            args.focus_view_points,
            args.modes,
            args.background_theme,
        )
        print(json.dumps({'room': room, **summary}, ensure_ascii=False))
        if args.render:
            render_assets(room_out, args.blender_bin, args.background_theme)
            grid = build_story_grid(room_out, room_tag, args.background_theme)
            if grid is not None:
                print(grid)


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(description='One-file delivery bundle for superpoint generation and Blender analysis views')
    sub = ap.add_subparsers(dest='cmd', required=True)

    p_gen = sub.add_parser('generate', help='Generate superpoint labels')
    p_gen.add_argument('--data_root', type=Path, default=DEFAULT_DATA_ROOT)
    p_gen.add_argument('--room', type=str, required=True)
    p_gen.add_argument('--output_root', type=Path, required=True)
    p_gen.add_argument('--method', choices=['segmentator', 'canonical'], default='segmentator')
    p_gen.add_argument('--voxel_size', type=float, default=0.12)
    p_gen.add_argument('--normal_bins', type=int, default=8)
    p_gen.add_argument('--knn_k', type=int, default=50)
    p_gen.add_argument('--k_thresh', type=float, default=0.01)
    p_gen.add_argument('--seg_min_verts', type=int, default=20)
    p_gen.add_argument('--segmentator_cpp', type=Path, default=DEFAULT_SEGMENTATOR_CPP)
    p_gen.add_argument('--segmentator_build_dir', type=Path, default=DEFAULT_SEGMENTATOR_BUILD_DIR)
    p_gen.add_argument('--build_verbose', action='store_true')
    p_gen.set_defaults(func=cmd_generate)

    p_story = sub.add_parser('story', help='Build Blender-ready superpoint analysis views')
    p_story.add_argument('--data_root', type=Path, default=DEFAULT_DATA_ROOT)
    p_story.add_argument('--rooms', nargs='+', required=True)
    p_story.add_argument('--output_root', type=Path, required=True)
    p_story.add_argument('--source', choices=['existing', 'canonical', 'segmentator'], default='segmentator')
    p_story.add_argument('--label_npy', type=Path, default=None)
    p_story.add_argument('--voxel_size', type=float, default=0.12)
    p_story.add_argument('--normal_bins', type=int, default=8)
    p_story.add_argument('--knn_k', type=int, default=50)
    p_story.add_argument('--k_thresh', type=float, default=0.01)
    p_story.add_argument('--seg_min_verts', type=int, default=20)
    p_story.add_argument('--segmentator_cpp', type=Path, default=DEFAULT_SEGMENTATOR_CPP)
    p_story.add_argument('--segmentator_build_dir', type=Path, default=DEFAULT_SEGMENTATOR_BUILD_DIR)
    p_story.add_argument('--build_verbose', action='store_true')
    p_story.add_argument('--max_points', type=int, default=38000)
    p_story.add_argument('--boundary_knn', type=int, default=20)
    p_story.add_argument('--sp_boundary_ratio', type=float, default=0.90)
    p_story.add_argument('--seg_boundary_ratio', type=float, default=0.30)
    p_story.add_argument('--focus_knn', type=int, default=26)
    p_story.add_argument('--focus_label', type=int, default=None)
    p_story.add_argument('--focus_view_points', type=int, default=12000)
    p_story.add_argument('--background_theme', choices=['white', 'light', 'dark'], default='white')
    p_story.add_argument('--modes', nargs='+', default=['semantic_family', 'semantic_family_cutaway', 'superpoint', 'superpoint_cutaway', 'sp_boundary', 'segment_boundary', 'boundary_relation', 'focus_superpoint'])
    p_story.add_argument('--render', action='store_true')
    p_story.add_argument('--blender_bin', type=Path, default=DEFAULT_BLENDER_BIN)
    p_story.set_defaults(func=cmd_story)
    return ap


def main() -> None:
    ap = build_parser()
    args = ap.parse_args()
    args.func(args)


if __name__ == '__main__':
    main()