| |
| 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() |
|
|
