#!/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()