#!/usr/bin/env python3 from __future__ import annotations import argparse import json from pathlib import Path import numpy as np from scipy.spatial import cKDTree def stable_hash_color(uid: int) -> np.ndarray: h = (int(uid) * 1103515245 + 12345) & 0x7FFFFFFF return np.array([(h >> 0) & 255, (h >> 8) & 255, (h >> 16) & 255], dtype=np.uint8) def sample_indices(n: int, max_points: int | None, rng: np.random.Generator, priority_masks: list[np.ndarray] | None = None) -> np.ndarray: if max_points is None or n <= max_points: return np.arange(n, dtype=np.int64) taken = np.zeros(n, dtype=bool) chosen = [] if priority_masks: for mask in priority_masks: idx = np.flatnonzero(mask & ~taken) if idx.size == 0: continue remain = max_points - (0 if not chosen else sum(len(x) for x in chosen)) if remain <= 0: break if idx.size > remain: idx = rng.choice(idx, size=remain, replace=False) taken[idx] = True chosen.append(np.asarray(idx, dtype=np.int64)) remain = max_points - (0 if not chosen else sum(len(x) for x in chosen)) if remain > 0: idx = np.flatnonzero(~taken) if idx.size > remain: idx = rng.choice(idx, size=remain, replace=False) chosen.append(np.asarray(idx, dtype=np.int64)) out = np.concatenate(chosen) if chosen else np.arange(max_points, dtype=np.int64) return np.sort(out.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(full_labels: np.ndarray, sampled_labels: np.ndarray, sampled_boundary: np.ndarray) -> int: uniq_s, counts_s = np.unique(sampled_labels, return_counts=True) b_uniq, b_counts = np.unique(sampled_labels[sampled_boundary], return_counts=True) b_map = {int(u): int(c) for u, c in zip(b_uniq, b_counts)} full_uniq, full_counts = np.unique(full_labels, return_counts=True) full_map = {int(u): int(c) for u, c in zip(full_uniq, full_counts)} boundary_counts = np.array([b_map.get(int(u), 0) for u in uniq_s], dtype=np.float64) full_sizes = np.array([full_map.get(int(u), 0) for u in uniq_s], dtype=np.float64) lo = max(80, np.percentile(full_sizes, 60)) hi = max(lo, np.percentile(full_sizes, 97)) score = (boundary_counts / np.maximum(counts_s, 1)) * np.log1p(full_sizes) cand = (full_sizes >= lo) & (full_sizes <= hi) & (boundary_counts > 0) if cand.any(): score = np.where(cand, score, -1.0) return int(uniq_s[int(np.argmax(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 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) -> np.ndarray: rgb = np.zeros((len(superpoint), 3), dtype=np.uint8) if mode == 'superpoint': for u in np.unique(superpoint): rgb[superpoint == u] = stable_hash_color(int(u)) return rgb if mode == 'sp_boundary': rgb[:] = np.array([56, 56, 60], dtype=np.uint8) rgb[sp_boundary] = np.array([255, 165, 72], dtype=np.uint8) return rgb if mode == 'segment_boundary': rgb[:] = np.array([56, 56, 60], dtype=np.uint8) rgb[seg_boundary] = np.array([74, 210, 255], dtype=np.uint8) return rgb if mode == 'boundary_relation': both = sp_boundary & seg_boundary sp_only = sp_boundary & ~seg_boundary seg_only = ~sp_boundary & seg_boundary rgb[:] = np.array([28, 28, 34], dtype=np.uint8) rgb[seg_only] = np.array([74, 170, 255], dtype=np.uint8) rgb[sp_only] = np.array([255, 185, 72], dtype=np.uint8) rgb[both] = np.array([255, 84, 84], dtype=np.uint8) return rgb if mode == 'focus_superpoint': focus_mask = superpoint == int(focus_label) same_seg = np.zeros_like(focus_mask) if np.any(focus_mask): dom_seg = int(np.bincount(segment[focus_mask]).argmax()) same_seg = segment == dom_seg rgb[:] = np.array([20, 20, 24], dtype=np.uint8) rgb[same_seg] = np.array([66, 82, 95], dtype=np.uint8) rgb[focus_shell] = np.array([102, 224, 245], dtype=np.uint8) rgb[focus_mask] = np.array([255, 108, 70], dtype=np.uint8) return rgb raise ValueError(mode) def save_asset(path: Path, coord: np.ndarray, rgb: np.ndarray, meta: dict): 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 main(): ap = argparse.ArgumentParser(description='Prepare Blender-ready superpoint story assets') here = Path(__file__).resolve().parent default_root = (here.parent.parent / '_work_biptv3' / 'pointcept_framework' / 'data' / 's3dis_official').resolve() ap.add_argument('--data_root', type=Path, default=default_root) ap.add_argument('--room', type=str, required=True) ap.add_argument('--label_npy', type=Path, default=None) ap.add_argument('--out_dir', type=Path, required=True) ap.add_argument('--max_points', type=int, default=45000) ap.add_argument('--boundary_knn', type=int, default=20) ap.add_argument('--sp_boundary_ratio', type=float, default=0.90) ap.add_argument('--seg_boundary_ratio', type=float, default=0.30) ap.add_argument('--focus_knn', type=int, default=26) ap.add_argument('--focus_label', type=int, default=None) ap.add_argument('--focus_view_points', type=int, default=12000) ap.add_argument('--modes', nargs='+', default=['superpoint', 'sp_boundary', 'segment_boundary', 'boundary_relation', 'focus_superpoint']) args = ap.parse_args() room_dir = args.data_root / args.room coord = np.load(room_dir / 'coord.npy').astype(np.float32) segment = np.load(room_dir / 'segment.npy').reshape(-1).astype(np.int64) label_path = args.label_npy or (room_dir / 'superpoint.npy') superpoint = np.load(label_path).reshape(-1).astype(np.int64) if not (len(coord) == len(segment) == len(superpoint)): raise ValueError('coord / segment / superpoint length mismatch') rng = np.random.default_rng(0) base_idx = sample_indices(len(coord), args.max_points, rng) coord_b = coord[base_idx] segment_b = segment[base_idx] superpoint_b = superpoint[base_idx] sp_boundary_b = knn_boundary_mask(coord_b, superpoint_b, k=args.boundary_knn, diff_ratio=args.sp_boundary_ratio) seg_boundary_b = knn_boundary_mask(coord_b, segment_b, k=args.boundary_knn, diff_ratio=args.seg_boundary_ratio) focus_label = int(args.focus_label) if args.focus_label is not None else auto_focus_label(superpoint, superpoint_b, sp_boundary_b) room_tag = args.room.replace('/', '_') summary = { 'room': args.room, 'label_npy': str(label_path), 'points_full': int(len(coord)), 'points_sampled': int(len(coord_b)), 'unique_superpoints_full': int(np.unique(superpoint).size), 'unique_superpoints_sampled': int(np.unique(superpoint_b).size), 'sp_boundary_points_sampled': int(sp_boundary_b.sum()), 'segment_boundary_points_sampled': int(seg_boundary_b.sum()), 'focus_label': int(focus_label), 'focus_label_points_full': int((superpoint == focus_label).sum()), 'focus_label_points_sampled': int((superpoint_b == focus_label).sum()), 'boundary_knn': int(args.boundary_knn), 'sp_boundary_ratio': float(args.sp_boundary_ratio), 'seg_boundary_ratio': float(args.seg_boundary_ratio), 'modes': list(args.modes), } args.out_dir.mkdir(parents=True, exist_ok=True) (args.out_dir / f'{room_tag}_summary.json').write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding='utf-8') for mode in args.modes: if mode == 'focus_superpoint': focus_mask_full = superpoint == focus_label focus_idx = local_focus_indices(coord, focus_mask_full, args.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=args.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) meta = dict(summary) meta.update({'mode': mode, 'background': 'dark', 'sampled_points': int(len(coord_m))}) save_asset(args.out_dir / f'{room_tag}_{mode}.npz', coord_m, rgb, meta) print('ASSET', args.out_dir / f'{room_tag}_{mode}.npz') continue rgb = build_rgb(mode, superpoint_b, segment_b, sp_boundary_b, seg_boundary_b, focus_label, np.zeros(len(coord_b), dtype=bool)) bg = 'light' if mode == 'superpoint' else 'dark' meta = dict(summary) meta.update({'mode': mode, 'background': bg, 'sampled_points': int(len(coord_b))}) save_asset(args.out_dir / f'{room_tag}_{mode}.npz', coord_b, rgb, meta) print('ASSET', args.out_dir / f'{room_tag}_{mode}.npz') if __name__ == '__main__': main()