"""Scene feature point cloud builder from raw HoHo dataset samples. V10 cache layout (additions over v7): - per-point RGB feature `scene_rgb` (uint8, 0..255): * COLMAP points use `point3D.color` from the COLMAP record * depth-unprojected points inherit the colour of their nearest COLMAP neighbour in raw-world space (no raw images are released) * camera tokens use a sentinel (0, 0, 0) - metric voxel downsampling across ALL candidate points (COLMAP + depth + camera) before the Tier-1/Tier-2 budget cap. At most one random point is kept per `voxel_size_m`-side voxel in raw world coordinates. Set `voxel_size_m <= 0` to disable. Inherited from v7: - robust normalisation (median centre + P95 scale on COLMAP+camera union) - bbox_R kept identity (random yaw is applied as data augmentation in the loader) - separate per-point geometric and semantic confidence channels - soft top-2 gestalt labels (`scene_gestalt_ids`/`scene_gestalt_id2` + `scene_gestalt_w1`) - 3x3 majority gestalt/ADE projection on COLMAP points - fixed split: 4096 COLMAP/camera tokens + 4096 depth-unprojected tokens - camera tokens live only in the COLMAP half (type=2, Tier 0) - Tier 1 structural Gestalt classes are preferred before weaker Tier 2 classes - Tier 2 points require an ADE house/foreground label - hard skip on empty COLMAP / missing cam metadata / ambiguous image matches """ import io from typing import Optional, Tuple import numpy as np from PIL import Image from hoho2025.example_solutions import _cam_matrix_from_image, get_fitted_dense_depth from color_mappings import ade20k_color_mapping, gestalt_color_mapping # ------------------------------------------------------------------ # Gestalt / ADE20K palette — packed int32 LUTs built once at import time. # ------------------------------------------------------------------ GESTALT_CLASSES = list(gestalt_color_mapping.keys()) # 28 classes, order = ID ADE20K_CLASSES = list(ade20k_color_mapping.keys()) def _build_palette_lut(color_mapping, start_id: int = 0): """Return (packed_sorted, ids_sorted) arrays for searchsorted lookup.""" colors = np.array([rgb for _, rgb in color_mapping.items()], dtype=np.int32) packed = (colors[:, 0] << 16) | (colors[:, 1] << 8) | colors[:, 2] ids = np.arange(start_id, start_id + len(packed), dtype=np.int64) order = np.argsort(packed) return packed[order], ids[order] _GESTALT_PACKED_S, _GESTALT_IDS_S = _build_palette_lut(gestalt_color_mapping, start_id=0) _ADE_PACKED_S, _ADE_IDS_S = _build_palette_lut(ade20k_color_mapping, start_id=1) def _sample_palette( img_np: np.ndarray, vi: np.ndarray, ui: np.ndarray, packed_sorted: np.ndarray, ids_sorted: np.ndarray, default: int = 0, ) -> np.ndarray: """Look up palette class IDs for pixels at (vi, ui) without converting the whole image.""" rgb = img_np[vi, ui].astype(np.int32) packed = (rgb[:, 0] << 16) | (rgb[:, 1] << 8) | rgb[:, 2] pos = np.searchsorted(packed_sorted, packed) in_range = pos < len(packed_sorted) hit = in_range & (packed_sorted[np.minimum(pos, len(packed_sorted) - 1)] == packed) return np.where(hit, ids_sorted[np.minimum(pos, len(ids_sorted) - 1)], default).astype(np.int64) def gestalt_img_to_ids(gest_np: np.ndarray) -> np.ndarray: H, W = gest_np.shape[:2] vi, ui = np.mgrid[0:H, 0:W] return _sample_palette( gest_np, vi.ravel(), ui.ravel(), _GESTALT_PACKED_S, _GESTALT_IDS_S, default=0, ).reshape(H, W) def ade20k_img_to_ids(ade_np: np.ndarray) -> np.ndarray: H, W = ade_np.shape[:2] vi, ui = np.mgrid[0:H, 0:W] return _sample_palette( ade_np, vi.ravel(), ui.ravel(), _ADE_PACKED_S, _ADE_IDS_S, default=0, ).reshape(H, W) def _ids_from_names(names: tuple[str, ...], all_names: list[str], start_id: int = 0) -> frozenset: return frozenset(start_id + all_names.index(name) for name in names if name in all_names) # Gestalt classes used in the fixed-budget scene cache. The user-facing # "flashing_end" and "transition line" names correspond to palette keys # "flashing_end_point" and "transition_line" in this codebase. GESTALT_TIER1_NAMES = ( "apex", "eave_end_point", "flashing_end_point", "eave", "ridge", "rake", "valley", ) GESTALT_TIER2_NAMES = ( "ground_line", "flashing", "step_flashing", "hip", "fascia", "transition_line", ) TIER1_GESTALT_IDS: frozenset = _ids_from_names(GESTALT_TIER1_NAMES, GESTALT_CLASSES) TIER2_GESTALT_IDS: frozenset = _ids_from_names(GESTALT_TIER2_NAMES, GESTALT_CLASSES) PRIORITY_GESTALT_IDS: frozenset = TIER1_GESTALT_IDS | TIER2_GESTALT_IDS _TIER1_GESTALT_ARR = np.array(sorted(TIER1_GESTALT_IDS), dtype=np.int64) _TIER2_GESTALT_ARR = np.array(sorted(TIER2_GESTALT_IDS), dtype=np.int64) _PRIORITY_GESTALT_ARR = np.array(sorted(PRIORITY_GESTALT_IDS), dtype=np.int64) ADE_HOUSE_FOREGROUND_NAMES = ( "wall", "building;edifice", "windowpane;window", "door;double;door", "house", "column;pillar", "skyscraper", "stairs;steps", ) ADE_HOUSE_FOREGROUND_IDS: frozenset = _ids_from_names( ADE_HOUSE_FOREGROUND_NAMES, ADE20K_CLASSES, start_id=1, ) _ADE_HOUSE_FOREGROUND_ARR = np.array(sorted(ADE_HOUSE_FOREGROUND_IDS), dtype=np.int64) def _ade_house_foreground_mask(ade_ids: np.ndarray) -> np.ndarray: return np.isin(ade_ids, _ADE_HOUSE_FOREGROUND_ARR) # ------------------------------------------------------------------ # Image decode + COLMAP plumbing # ------------------------------------------------------------------ def _image_to_rgb_array(img_like) -> np.ndarray: if isinstance(img_like, np.ndarray): arr = img_like elif isinstance(img_like, dict) and img_like.get("bytes") is not None: arr = np.array(Image.open(io.BytesIO(img_like["bytes"])).convert("RGB")) elif hasattr(img_like, "convert"): arr = np.array(img_like.convert("RGB")) else: arr = np.array(img_like) if arr.ndim == 2: arr = np.repeat(arr[..., None], 3, axis=-1) if arr.shape[-1] > 3: arr = arr[..., :3] return arr.astype(np.uint8, copy=False) def _image_to_pil(img_like): """Decode image-like dataset values while preserving depth image modes.""" if isinstance(img_like, Image.Image): return img_like if isinstance(img_like, dict) and img_like.get("bytes") is not None: img = Image.open(io.BytesIO(img_like["bytes"])) img.load() return img if isinstance(img_like, np.ndarray): return Image.fromarray(img_like) return img_like def _camera_for_image(colmap_rec, col_img): try: return col_img.camera except AttributeError: return colmap_rec.cameras[col_img.camera_id] def _observed_point_mask(col_img, point_ids: np.ndarray) -> np.ndarray: ids_attr = getattr(col_img, "point3D_ids", None) if ids_attr is not None: ids = ids_attr() if callable(ids_attr) else ids_attr ids = np.asarray(ids, dtype=np.int64) ids = ids[ids >= 0] if ids.size > 0: return np.isin(point_ids, ids) if hasattr(col_img, "has_point3D"): observed = np.fromiter( (bool(col_img.has_point3D(int(pid))) for pid in point_ids), dtype=bool, count=len(point_ids), ) if observed.any(): return observed return np.ones(len(point_ids), dtype=bool) # ------------------------------------------------------------------ # COLMAP point + confidence extraction # ------------------------------------------------------------------ def colmap_points_xyz(colmap_rec) -> np.ndarray: """Extract (N, 3) float32 world-space COLMAP points (xyz only).""" if not colmap_rec.points3D: return np.zeros((0, 3), dtype=np.float32) return np.array([p.xyz for p in colmap_rec.points3D.values()], dtype=np.float32) def colmap_points_xyz_ids(colmap_rec): if not colmap_rec.points3D: return np.zeros(0, dtype=np.int64), np.zeros((0, 3), dtype=np.float32) ids, xyz = [], [] for pid, p in colmap_rec.points3D.items(): ids.append(int(pid)) xyz.append(p.xyz) return np.asarray(ids, dtype=np.int64), np.asarray(xyz, dtype=np.float32) def colmap_points_full(colmap_rec): """Return (ids, xyz, track_len, reproj_err, rgb) for all COLMAP 3D points. `rgb` is uint8 (N, 3) sourced from each point's COLMAP `color` field. """ if not colmap_rec.points3D: return ( np.zeros(0, dtype=np.int64), np.zeros((0, 3), dtype=np.float32), np.zeros(0, dtype=np.float32), np.zeros(0, dtype=np.float32), np.zeros((0, 3), dtype=np.uint8), ) ids, xyz, tlen, err, rgb = [], [], [], [], [] for pid, p in colmap_rec.points3D.items(): ids.append(int(pid)) xyz.append(p.xyz) try: tlen.append(int(p.track.length())) except Exception: try: tlen.append(int(len(p.track.elements))) except Exception: tlen.append(2) try: err.append(float(p.error)) except Exception: err.append(1.0) try: rgb.append(np.asarray(p.color, dtype=np.uint8)) except Exception: rgb.append(np.zeros(3, dtype=np.uint8)) return ( np.asarray(ids, dtype=np.int64), np.asarray(xyz, dtype=np.float32), np.asarray(tlen, dtype=np.float32), np.asarray(err, dtype=np.float32), np.asarray(rgb, dtype=np.uint8), ) def voxel_downsample_indices( xyz: np.ndarray, voxel_size: float, rng: np.random.Generator, ) -> np.ndarray: """Return indices into `xyz` keeping one random point per metric voxel. Voxel edges are aligned to a grid of side `voxel_size` in the same units as `xyz` (raw COLMAP world space). When `voxel_size <= 0` the function returns a permutation of all indices (no-op). """ n = len(xyz) if n == 0 or voxel_size <= 0: return np.arange(n, dtype=np.int64) keys = np.floor(xyz / voxel_size).astype(np.int64) # lexsort by (z, y, x) so groups of identical voxel keys are contiguous order = np.lexsort(keys.T[::-1]) keys_sorted = keys[order] diff = np.any(keys_sorted[1:] != keys_sorted[:-1], axis=1) starts = np.concatenate([[0], np.where(diff)[0] + 1]) ends = np.concatenate([starts[1:], [n]]) sizes = ends - starts offsets = rng.integers(0, sizes) # one random pick per voxel return order[starts + offsets].astype(np.int64, copy=False) def nearest_colmap_rgb( query_xyz: np.ndarray, colmap_xyz: np.ndarray, colmap_rgb: np.ndarray, ) -> np.ndarray: """For each query point, return the RGB of its nearest COLMAP neighbour. Falls back to all-zero RGB when COLMAP is empty. Uses scipy cKDTree when available; otherwise a chunked brute-force search. """ n_q = len(query_xyz) if n_q == 0: return np.zeros((0, 3), dtype=np.uint8) if len(colmap_xyz) == 0: return np.zeros((n_q, 3), dtype=np.uint8) try: from scipy.spatial import cKDTree tree = cKDTree(colmap_xyz.astype(np.float32, copy=False)) _, idx = tree.query(query_xyz.astype(np.float32, copy=False), k=1) idx = np.asarray(idx, dtype=np.int64).reshape(-1) except Exception: # Brute-force fallback in float32 chunks to bound memory. idx = np.zeros(n_q, dtype=np.int64) c = colmap_xyz.astype(np.float32, copy=False) chunk = 4096 for s in range(0, n_q, chunk): q = query_xyz[s : s + chunk].astype(np.float32, copy=False) d2 = ((q[:, None, :] - c[None, :, :]) ** 2).sum(axis=-1) idx[s : s + chunk] = d2.argmin(axis=1) return colmap_rgb[idx] def colmap_camera_centers(colmap_rec) -> np.ndarray: centers = [] for img in colmap_rec.images.values(): R, t = _cam_matrix_from_image(img) centers.append(R.T @ (-t)) if not centers: return np.zeros((0, 3), dtype=np.float32) return np.array(centers, dtype=np.float32) # ------------------------------------------------------------------ # Robust normalisation (no yaw alignment — that's a training augmentation now) # ------------------------------------------------------------------ def _robust_norm_params(colmap_xyz: np.ndarray, cam_centers: np.ndarray): """Return (center, scale) for the scene. Rotation is identity by design.""" pts = [] if colmap_xyz.size: pts.append(colmap_xyz) if cam_centers.size: pts.append(cam_centers) pts = np.concatenate(pts, axis=0).astype(np.float64) center = np.median(pts, axis=0).astype(np.float32) d = np.linalg.norm(pts - center, axis=1) scale = float(max(np.percentile(d, 95.0), 1e-3)) return center, scale def _normalise(pts: np.ndarray, center: np.ndarray, scale: float) -> np.ndarray: return ((pts - center) / scale).astype(np.float32) # ------------------------------------------------------------------ # Multi-view semantic projection (3x3 majority + soft top-2 + sem confidence) # ------------------------------------------------------------------ def _project_to_pixel(pts: np.ndarray, col_img, cam): """Project (N, 3) world points to pixel space of a single camera.""" cam_w = getattr(cam, "width", 0) or 0 cam_h = getattr(cam, "height", 0) or 0 if cam_w <= 0 or cam_h <= 0: return None R, t = _cam_matrix_from_image(col_img) p_cam = pts.astype(np.float64) @ R.T + t[None] z = p_cam[:, 2] in_front = z > 1e-8 if not in_front.any(): return None K = cam.calibration_matrix() u_proj = p_cam[:, 0] / np.maximum(z, 1e-8) * K[0, 0] + K[0, 2] v_proj = p_cam[:, 1] / np.maximum(z, 1e-8) * K[1, 1] + K[1, 2] return u_proj, v_proj, in_front, float(cam_w), float(cam_h) def _gather_3x3_votes( img_np: np.ndarray, vi: np.ndarray, ui: np.ndarray, packed_sorted: np.ndarray, ids_sorted: np.ndarray, n_classes: int, default_id: int, ) -> np.ndarray: """Sample 3x3 neighbourhoods and return per-point class vote counts.""" H, W = img_np.shape[:2] N = len(vi) counts = np.zeros((N, n_classes), dtype=np.uint16) for dv in (-1, 0, 1): for du in (-1, 0, 1): v2 = vi + dv u2 = ui + du ok = (v2 >= 0) & (v2 < H) & (u2 >= 0) & (u2 < W) if not ok.any(): continue idx = np.where(ok)[0] cls = _sample_palette( img_np, v2[ok], u2[ok], packed_sorted, ids_sorted, default=default_id, ) np.add.at(counts, (idx, cls), 1) return counts def _resolve_colmap_image(colmap_by_name: dict, img_id: str): """Return the COLMAP image whose name matches img_id.""" exact = colmap_by_name.get(img_id) if exact is not None: return exact matches = [v for k, v in colmap_by_name.items() if img_id in k] if len(matches) == 1: return matches[0] return None def project_semantics_to_colmap_points( sample: dict, colmap_rec, point_ids: np.ndarray, point_xyz: np.ndarray, ) -> dict: """Vote projected Gestalt/ADE labels onto selected COLMAP 3D points (3x3).""" n_points = len(point_xyz) n_g = len(GESTALT_CLASSES) n_a = len(ADE20K_CLASSES) + 1 gestalt_counts = np.zeros((n_points, n_g), dtype=np.uint32) ade_counts = np.zeros((n_points, n_a), dtype=np.uint32) image_ids = sample.get("image_ids", []) gestalt_imgs = sample.get("gestalt", []) ade_imgs = sample.get("ade", []) colmap_by_name = {col_img.name: col_img for col_img in colmap_rec.images.values()} for i, img_id in enumerate(image_ids): col_img = _resolve_colmap_image(colmap_by_name, img_id) if col_img is None: continue cam = _camera_for_image(colmap_rec, col_img) observed_mask = _observed_point_mask(col_img, point_ids) obs_idx = np.nonzero(observed_mask)[0] if obs_idx.size == 0: continue proj = _project_to_pixel(point_xyz[obs_idx], col_img, cam) if proj is None: continue u_proj, v_proj, in_front, cam_w, cam_h = proj def _valid_pixels(H: int, W: int): u = np.rint(u_proj * (W / cam_w)).astype(np.int64) v = np.rint(v_proj * (H / cam_h)).astype(np.int64) ok = in_front & (u >= 0) & (u < W) & (v >= 0) & (v < H) return u, v, ok if i < len(gestalt_imgs) and gestalt_imgs[i] is not None: gest_np = _image_to_rgb_array(gestalt_imgs[i]) ui, vi, ok = _valid_pixels(*gest_np.shape[:2]) if ok.any(): votes = _gather_3x3_votes( gest_np, vi[ok], ui[ok], _GESTALT_PACKED_S, _GESTALT_IDS_S, n_classes=n_g, default_id=0, ) rows = obs_idx[ok] gestalt_counts[rows] += votes.astype(np.uint32) if i < len(ade_imgs) and ade_imgs[i] is not None: ade_np = _image_to_rgb_array(ade_imgs[i]) ui, vi, ok = _valid_pixels(*ade_np.shape[:2]) if ok.any(): votes = _gather_3x3_votes( ade_np, vi[ok], ui[ok], _ADE_PACKED_S, _ADE_IDS_S, n_classes=n_a, default_id=0, ) rows = obs_idx[ok] ade_counts[rows] += votes.astype(np.uint32) gest_total = gestalt_counts.sum(axis=1) has_vote = gest_total > 0 top1 = np.argmax(gestalt_counts, axis=1) counts2 = gestalt_counts.copy() counts2[np.arange(n_points), top1] = 0 top2 = np.argmax(counts2, axis=1) has_second = counts2[np.arange(n_points), top2] > 0 gest_id1 = np.where(has_vote, top1, -1).astype(np.int64) gest_id2 = np.where(has_vote & has_second, top2, -1).astype(np.int64) top1_count = gestalt_counts[np.arange(n_points), top1].astype(np.float32) gest_w1 = np.where( has_vote, top1_count / np.maximum(gest_total.astype(np.float32), 1.0), 1.0, ).astype(np.float32) ade_ids = ade_counts.argmax(axis=1).astype(np.int64) return { "gestalt_id1": gest_id1, "gestalt_id2": gest_id2, "gestalt_w1": gest_w1, "ade_ids": ade_ids, "gest_top1_count": top1_count, } # ------------------------------------------------------------------ # Tier-targeted depth unprojection (single-pixel labels) # ------------------------------------------------------------------ def fit_depth_label_cache( sample: dict, colmap_rec, ) -> list[dict]: """Per-image preprocessing cache shared by stage-1 and stage-2 builders. For each image with valid depth fitting, returns the fitted depth map plus the gestalt + ADE label images resized to the depth resolution and the K/R/t needed to unproject pixels into world space. Stage 1's ``unproject_depth_tiered`` and stage 2's hull-cropped depth extractor both consume this cache, avoiding a second pass through ``get_fitted_dense_depth`` (Metric3Dv2) and a second PIL decode per image. """ K_mats = np.array(sample["K"]) R_mats = np.array(sample["R"]) t_vecs = np.array(sample["t"]) if t_vecs.ndim == 3: t_vecs = t_vecs[:, :, 0] pose_only = list(sample.get("pose_only_in_colmap", [False] * len(sample["image_ids"]))) out: list[dict] = [] for i, (depth_img, gest_img, ade_img, img_id) in enumerate( zip(sample["depth"], sample["gestalt"], sample["ade"], sample["image_ids"]) ): if pose_only[i] or depth_img is None: continue Ki = K_mats[i] if Ki.shape != (3, 3) or Ki[0, 0] == 0: continue Ri = R_mats[i] ti = t_vecs[i] if not (np.isfinite(Ki).all() and np.isfinite(Ri).all() and np.isfinite(ti).all()): continue fx, fy = float(Ki[0, 0]), float(Ki[1, 1]) if abs(fx) < 1e-6 or abs(fy) < 1e-6: continue depth_pil = _image_to_pil(depth_img) ade_pil = _image_to_pil(ade_img) if ade_img is not None else None gest_pil = _image_to_pil(gest_img) if gest_img is not None else None try: depth_fitted, _, found, _, _ = get_fitted_dense_depth( depth_pil, colmap_rec, img_id, ade_pil, verbose=False, ) except Exception: continue if not found: continue H, W = depth_fitted.shape gest_np = (_image_to_rgb_array(gest_pil.resize((W, H), Image.NEAREST)) if gest_pil is not None else None) ade_np = (_image_to_rgb_array(ade_pil.resize((W, H), Image.NEAREST)) if ade_pil is not None else None) out.append({ "i": int(i), "img_id": img_id, "depth_fitted": depth_fitted, "gest_np": gest_np, "ade_np": ade_np, "Ki": Ki.astype(np.float32, copy=False), "Ri": Ri.astype(np.float32, copy=False), "ti": ti.astype(np.float32, copy=False), }) return out def unproject_depth_tiered( sample: dict, colmap_rec, n_per_image: int = 4096, rng: Optional[np.random.Generator] = None, depth_cache: Optional[list[dict]] = None, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Unproject Tier 1/2 Gestalt pixels. Returns (xyz, gestalt_ids, ade_ids, conf). Uses a single-pixel sample (no 3x3 voting): the pixel's own Gestalt/ADE IDs become the point labels. Tier 2 pixels are accepted only where ADE maps to the house/foreground subset, mirroring COLMAP token selection. When ``depth_cache`` (built by :func:`fit_depth_label_cache`) is supplied, the per-image PIL decode + Metric3Dv2 fitting is skipped — entries are consumed directly. Otherwise we build the cache implicitly. """ if rng is None: rng = np.random.default_rng() if depth_cache is None: depth_cache = fit_depth_label_cache(sample, colmap_rec) all_xyz, all_gids, all_ade, all_conf = [], [], [], [] for entry in depth_cache: if entry.get("gest_np") is None or entry.get("ade_np") is None: # Stage 1's tier filter needs both gestalt and ADE labels per pixel. continue depth_fitted = entry["depth_fitted"] gest_np = entry["gest_np"] ade_np = entry["ade_np"] Ki = entry["Ki"] Ri = entry["Ri"] ti = entry["ti"] H, W = depth_fitted.shape gid_map = gestalt_img_to_ids(gest_np) ade_id_map = ade20k_img_to_ids(ade_np) valid_depth = depth_fitted > 0.1 tier1_mask = np.isin(gid_map, _TIER1_GESTALT_ARR) & valid_depth tier2_mask = ( np.isin(gid_map, _TIER2_GESTALT_ARR) & valid_depth & _ade_house_foreground_mask(ade_id_map) ) def _sample_flat(mask: np.ndarray, budget: int) -> np.ndarray: if budget <= 0: return np.zeros(0, dtype=np.int64) flat = np.flatnonzero(mask.ravel()) if flat.size > budget: flat = rng.choice(flat, budget, replace=False) return flat.astype(np.int64, copy=False) tier1_flat = _sample_flat(tier1_mask, n_per_image) tier2_flat = _sample_flat(tier2_mask, n_per_image - len(tier1_flat)) if tier1_flat.size == 0 and tier2_flat.size == 0: continue sel = np.concatenate([tier1_flat, tier2_flat], axis=0) ys = (sel // W).astype(np.int64) xs = (sel % W).astype(np.int64) depths = depth_fitted[ys, xs] fx, fy = Ki[0, 0], Ki[1, 1] cx, cy = Ki[0, 2], Ki[1, 2] x_cam = ((xs - cx) / fx) * depths y_cam = ((ys - cy) / fy) * depths z_cam = depths pts_cam = np.stack([x_cam, y_cam, z_cam], axis=1) pts_world = (pts_cam - ti[None]) @ Ri gids = gid_map[ys, xs].astype(np.int64) ade_ids_arr = ade_id_map[ys, xs].astype(np.int64) conf = np.clip(2.0 / np.maximum(depths, 0.5), 0.0, 1.0).astype(np.float32) all_xyz.append(pts_world.astype(np.float32)) all_gids.append(gids) all_ade.append(ade_ids_arr.astype(np.int64)) all_conf.append(conf) if not all_xyz: return ( np.zeros((0, 3), dtype=np.float32), np.zeros(0, dtype=np.int64), np.zeros(0, dtype=np.int64), np.zeros(0, dtype=np.float32), ) return ( np.concatenate(all_xyz, axis=0), np.concatenate(all_gids, axis=0), np.concatenate(all_ade, axis=0), np.concatenate(all_conf, axis=0), ) def unproject_depth_priority( sample: dict, colmap_rec, n_per_image: int = 4096, rng: Optional[np.random.Generator] = None, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Backward-compatible alias for the tiered depth sampler.""" return unproject_depth_tiered(sample, colmap_rec, n_per_image=n_per_image, rng=rng) # ------------------------------------------------------------------ # Main scene-input builder # ------------------------------------------------------------------ def _colmap_geom_conf(track_len: np.ndarray, reproj_err: np.ndarray) -> np.ndarray: """Geometric confidence of a COLMAP point in [0, 1].""" tl = np.clip(track_len / 5.0, 0.0, 1.0) er = np.exp(-np.maximum(reproj_err, 0.0) / 2.0) return (tl * er).astype(np.float32) def build_scene_input( sample: dict, n_pts: int = 8192, use_depth: bool = True, n_depth_per_image: int = 4096, voxel_size_m: float = 0.0, rng: Optional[np.random.Generator] = None, colmap_sem: Optional[dict] = None, depth_cache: Optional[list[dict]] = None, return_cache: bool = False, ) -> dict: """Build a fixed-size scene feature point cloud from one dataset sample. `n_pts` is split in half by provenance (default 8192 = 4096 + 4096): the first half is COLMAP/camera tokens, the second half is depth-unprojected tokens. If a scene cannot supply enough points for either half, that half is padded by resampling from the points already selected for the same half. `voxel_size_m` (metres in raw world space) thins the candidate pool before the budget cap: COLMAP, depth and camera candidates are jointly voxelised and at most one random point per voxel survives. Set to 0 to disable. Tier order: COLMAP half: T0 camera centres (type=2) T1 apex/eave_end/flashing_end/eave/ridge/rake/valley T2 ground_line/flashing/step_flashing/hip/fascia/transition_line only when ADE is house/foreground Depth half: T1 same structural Gestalt classes (type=1) T2 same weak Gestalt classes with house/foreground ADE """ if rng is None: rng = np.random.default_rng() colmap_rec = sample["colmap"] colmap_pids, colmap_xyz, colmap_tlen, colmap_err, colmap_rgb = colmap_points_full(colmap_rec) cam_centers = colmap_camera_centers(colmap_rec) if len(colmap_xyz) == 0: raise ValueError("scene has no COLMAP 3D points — cannot build scene input") bbox_center, bbox_scale = _robust_norm_params(colmap_xyz, cam_centers) bbox_R = np.eye(3, dtype=np.float32) # identity by design (yaw is augmentation) def norm(pts: np.ndarray) -> np.ndarray: return _normalise(pts, bbox_center, bbox_scale) # ---- COLMAP semantic projection ---------------------------------------- # Reuse a precomputed semantic projection when one is supplied — the # 3x3-voting pass is the dominant per-sample cost in stage-1 preprocess. if colmap_sem is None: sem = project_semantics_to_colmap_points(sample, colmap_rec, colmap_pids, colmap_xyz) else: sem = colmap_sem all_g1 = sem["gestalt_id1"] all_g2 = sem["gestalt_id2"] all_w1 = sem["gestalt_w1"] all_ade = sem["ade_ids"] top1_count = sem["gest_top1_count"] geom_conf_all = _colmap_geom_conf(colmap_tlen, colmap_err) sem_conf_all = np.clip(top1_count / 5.0, 0.0, 1.0).astype(np.float32) colmap_budget = n_pts // 2 depth_budget = n_pts - colmap_budget tier1_mask = np.isin(all_g1, _TIER1_GESTALT_ARR) tier2_mask = np.isin(all_g1, _TIER2_GESTALT_ARR) & _ade_house_foreground_mask(all_ade) # ---- Depth unprojection (once) + nearest-COLMAP RGB -------------------- if use_depth: # Build the per-image depth cache up here (instead of letting # unproject_depth_tiered build it internally) so we can capture it # for `return_cache=True` — otherwise stage 2 would redo Metric3Dv2 # depth fitting per image, defeating the shared-compute refactor. if depth_cache is None: depth_cache = fit_depth_label_cache(sample, colmap_rec) d_xyz, d_gids, d_ade, d_conf = unproject_depth_tiered( sample, colmap_rec, n_per_image=n_depth_per_image, rng=rng, depth_cache=depth_cache, ) else: d_xyz = np.zeros((0, 3), dtype=np.float32) d_gids = np.zeros(0, dtype=np.int64) d_ade = np.zeros(0, dtype=np.int64) d_conf = np.zeros(0, dtype=np.float32) # Build the COLMAP KDTree once; only query for depth points that survive # tier + voxel + budget selection. Avoids ~10-20x wasted queries on the # unprojected depth pool (which is typically ~80k while the final depth # half is ~4k). _colmap_kd = None if len(d_xyz) > 0 and len(colmap_xyz) > 0: try: from scipy.spatial import cKDTree _colmap_kd = cKDTree(colmap_xyz.astype(np.float32, copy=False)) except Exception: _colmap_kd = None def _depth_rgb_for(sel: np.ndarray) -> np.ndarray: if len(sel) == 0: return np.zeros((0, 3), dtype=np.uint8) if _colmap_kd is None or len(colmap_rgb) == 0: return np.zeros((len(sel), 3), dtype=np.uint8) _, idx = _colmap_kd.query(d_xyz[sel].astype(np.float32, copy=False), k=1) return colmap_rgb[np.asarray(idx, dtype=np.int64).reshape(-1)] cam_rgb = np.zeros((len(cam_centers), 3), dtype=np.uint8) # ---- Joint voxel downsampling ----------------------------------------- # Voxelise COLMAP + depth + camera candidates together in raw world space. # The kept-mask is sliced back per provenance and intersected with tier # masks below, so each Tier-1/Tier-2 selection picks from voxel survivors. if voxel_size_m and voxel_size_m > 0.0 and (len(colmap_xyz) + len(d_xyz) + len(cam_centers)) > 0: union_xyz = np.concatenate( [colmap_xyz, d_xyz, cam_centers.astype(np.float32, copy=False)], axis=0 ) keep_global = voxel_downsample_indices(union_xyz, float(voxel_size_m), rng) kept_flags = np.zeros(len(union_xyz), dtype=bool) kept_flags[keep_global] = True n_c = len(colmap_xyz) n_d = len(d_xyz) colmap_voxel_keep = kept_flags[:n_c] depth_voxel_keep = kept_flags[n_c : n_c + n_d] cam_voxel_keep = kept_flags[n_c + n_d :] else: colmap_voxel_keep = np.ones(len(colmap_xyz), dtype=bool) depth_voxel_keep = np.ones(len(d_xyz), dtype=bool) cam_voxel_keep = np.ones(len(cam_centers), dtype=bool) def _pick(mask: np.ndarray, budget: int) -> np.ndarray: if budget <= 0: return np.zeros(0, dtype=np.int64) idx = np.flatnonzero(mask) if len(idx) > budget: idx = rng.choice(idx, budget, replace=False) return idx.astype(np.int64, copy=False) def _empty_arrays(n: int = 0) -> tuple: return ( np.zeros((n, 3), dtype=np.float32), # 0: xyz np.zeros(n, dtype=np.int64), # 1: type_id np.full(n, -1, dtype=np.int64), # 2: gestalt_id1 np.full(n, -1, dtype=np.int64), # 3: gestalt_id2 np.ones(n, dtype=np.float32), # 4: gestalt_w1 np.zeros(n, dtype=np.int64), # 5: ade_id np.zeros(n, dtype=np.float32), # 6: geom_conf np.zeros(n, dtype=np.float32), # 7: sem_conf np.zeros((n, 3), dtype=np.uint8), # 8: rgb ) _N_FIELDS = 9 def _concat(parts: list[tuple]) -> tuple: if not parts: return _empty_arrays(0) return tuple(np.concatenate([p[i] for p in parts], axis=0) for i in range(_N_FIELDS)) def _take(arrs: tuple, idx: np.ndarray) -> tuple: return tuple(a[idx] for a in arrs) def _pad_or_trim(arrs: tuple, budget: int, fallback: Optional[tuple] = None) -> tuple: if budget <= 0: return _empty_arrays(0) n = len(arrs[0]) if n == 0 and fallback is not None and len(fallback[0]) > 0: n_fallback = len(fallback[0]) idx = rng.choice(n_fallback, budget, replace=n_fallback < budget) return _take(fallback, idx) if n == 0: return _empty_arrays(budget) if n > budget: idx = rng.choice(n, budget, replace=False) return _take(arrs, idx) if n < budget: pad = rng.choice(n, budget - n, replace=True) return tuple(np.concatenate([a, a[pad]], axis=0) for a in arrs) return arrs def _camera_arrays(sel: np.ndarray) -> tuple: n = len(sel) return ( norm(cam_centers[sel]), np.full(n, 2, dtype=np.int64), np.full(n, -1, dtype=np.int64), np.full(n, -1, dtype=np.int64), np.ones(n, dtype=np.float32), np.zeros(n, dtype=np.int64), np.ones(n, dtype=np.float32), np.ones(n, dtype=np.float32), cam_rgb[sel], ) def _colmap_arrays(sel: np.ndarray) -> tuple: n = len(sel) if n == 0: return _empty_arrays(0) return ( norm(colmap_xyz[sel]), np.zeros(n, dtype=np.int64), all_g1[sel], all_g2[sel], all_w1[sel], all_ade[sel], geom_conf_all[sel], sem_conf_all[sel], colmap_rgb[sel], ) def _depth_arrays(sel: np.ndarray) -> tuple: n = len(sel) if n == 0: return _empty_arrays(0) return ( norm(d_xyz[sel]), np.full(n, 1, dtype=np.int64), d_gids[sel], np.full(n, -1, dtype=np.int64), np.ones(n, dtype=np.float32), d_ade[sel], d_conf[sel], np.ones(n, dtype=np.float32), _depth_rgb_for(sel), ) # ---- COLMAP half: T0 cameras, then Tier 1, then Tier 2, then random fill col_parts: list[tuple] = [] remaining_col = colmap_budget picked_colmap = np.zeros(len(colmap_xyz), dtype=bool) cam_eligible = np.flatnonzero(cam_voxel_keep) if remaining_col > 0 and len(cam_eligible) > 0: n_cam = min(len(cam_eligible), remaining_col) sel = ( rng.choice(cam_eligible, n_cam, replace=False) if len(cam_eligible) > n_cam else cam_eligible ) col_parts.append(_camera_arrays(sel.astype(np.int64, copy=False))) remaining_col -= n_cam if remaining_col > 0: sel = _pick(tier1_mask & colmap_voxel_keep & ~picked_colmap, remaining_col) col_parts.append(_colmap_arrays(sel)) picked_colmap[sel] = True remaining_col -= len(sel) if remaining_col > 0: sel = _pick(tier2_mask & colmap_voxel_keep & ~picked_colmap, remaining_col) col_parts.append(_colmap_arrays(sel)) picked_colmap[sel] = True remaining_col -= len(sel) # Fill whatever is left from ordinary COLMAP points (still voxel-thinned). if remaining_col > 0: sel = _pick(colmap_voxel_keep & ~picked_colmap, remaining_col) col_parts.append(_colmap_arrays(sel)) picked_colmap[sel] = True remaining_col -= len(sel) col_arrays = _pad_or_trim(_concat(col_parts), colmap_budget) # ---- Depth half: Tier 1, then Tier 2 ------------------------------------ depth_parts: list[tuple] = [] if depth_budget > 0 and use_depth and len(d_xyz) > 0: d_tier1 = np.isin(d_gids, _TIER1_GESTALT_ARR) & depth_voxel_keep d_tier2 = ( np.isin(d_gids, _TIER2_GESTALT_ARR) & _ade_house_foreground_mask(d_ade) & depth_voxel_keep ) remaining_depth = depth_budget sel = _pick(d_tier1, remaining_depth) depth_parts.append(_depth_arrays(sel)) remaining_depth -= len(sel) if remaining_depth > 0: sel = _pick(d_tier2, remaining_depth) depth_parts.append(_depth_arrays(sel)) depth_arrays = _pad_or_trim(_concat(depth_parts), depth_budget, fallback=col_arrays) all_xyz = np.concatenate([col_arrays[0], depth_arrays[0]], axis=0) all_type = np.concatenate([col_arrays[1], depth_arrays[1]], axis=0) all_g1 = np.concatenate([col_arrays[2], depth_arrays[2]], axis=0) all_g2 = np.concatenate([col_arrays[3], depth_arrays[3]], axis=0) all_w1 = np.concatenate([col_arrays[4], depth_arrays[4]], axis=0) all_ade = np.concatenate([col_arrays[5], depth_arrays[5]], axis=0) all_geom_conf = np.concatenate([col_arrays[6], depth_arrays[6]], axis=0) all_sem_conf = np.concatenate([col_arrays[7], depth_arrays[7]], axis=0) all_rgb = np.concatenate([col_arrays[8], depth_arrays[8]], axis=0) out = { "scene_xyz": all_xyz, "scene_type_ids": all_type, "scene_gestalt_ids": all_g1, "scene_gestalt_id2": all_g2, "scene_gestalt_w1": all_w1, "scene_ade_ids": all_ade, "scene_geom_conf": all_geom_conf.astype(np.float32), "scene_sem_conf": all_sem_conf.astype(np.float32), "scene_rgb": all_rgb.astype(np.uint8), "bbox_center": bbox_center.astype(np.float32), "bbox_scale": float(bbox_scale), "bbox_R": bbox_R, } if return_cache: out["_cache"] = { "colmap": { "pids": colmap_pids, "xyz": colmap_xyz, "rgb": colmap_rgb, "track_len": colmap_tlen, "reproj_err": colmap_err, "sem": sem, }, # Either the cache passed in, or the one built above for use_depth. # None only when use_depth=False (then stage 2 has no depth to share). "depth": depth_cache, } return out # ------------------------------------------------------------------ # GT vertex builder # ------------------------------------------------------------------ # ------------------------------------------------------------------ # V11 chimney filter — drop small connected components from GT # ------------------------------------------------------------------ def _polygon_xy_area(verts_xy: np.ndarray) -> float: """Shoelace area of the convex polygon formed by 1..4 XY points. Vertices are first sorted by polar angle around their centroid so the polygon is non-self-intersecting. For < 3 points the area is zero. """ n = len(verts_xy) if n < 3: return 0.0 xy = verts_xy.astype(np.float64, copy=False) c = xy.mean(axis=0) ang = np.arctan2(xy[:, 1] - c[1], xy[:, 0] - c[0]) order = np.argsort(ang) p = xy[order] x, y = p[:, 0], p[:, 1] return 0.5 * float(abs(np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1)))) def filter_small_components( wf_vertices: np.ndarray, wf_edges: np.ndarray, max_verts_per_cc: int = 4, max_area_m2: float = 5.0, ) -> Tuple[np.ndarray, np.ndarray]: """Mark vertices/edges belonging to small connected components for removal. Operates on RAW (un-normalised, metric) coordinates. A connected component is removed iff it has at most ``max_verts_per_cc`` vertices AND its XY footprint area is strictly less than ``max_area_m2`` square metres. Args: wf_vertices: (N, 3) world-space GT vertices. wf_edges: (M, 2) edges as vertex-index pairs. Returns: (keep_v, keep_e): boolean masks of length N and M respectively. True means "retain" (not a chimney). """ n = int(len(wf_vertices)) m = int(len(wf_edges)) if n == 0: return np.ones(0, dtype=bool), np.ones(m, dtype=bool) parent = np.arange(n, dtype=np.int64) def find(i: int) -> int: while parent[i] != i: parent[i] = parent[parent[i]] i = parent[i] return i if m > 0: edges_int = np.asarray(wf_edges, dtype=np.int64) valid_e = ( (edges_int[:, 0] >= 0) & (edges_int[:, 1] >= 0) & (edges_int[:, 0] < n) & (edges_int[:, 1] < n) ) for a, b in edges_int[valid_e]: ra, rb = find(int(a)), find(int(b)) if ra != rb: parent[ra] = rb else: edges_int = np.zeros((0, 2), dtype=np.int64) valid_e = np.zeros(0, dtype=bool) roots = np.array([find(i) for i in range(n)], dtype=np.int64) keep_v = np.ones(n, dtype=bool) pts = np.asarray(wf_vertices, dtype=np.float32) for root in np.unique(roots): idx = np.where(roots == root)[0] if len(idx) <= max_verts_per_cc: area = _polygon_xy_area(pts[idx, :2]) if area < max_area_m2: keep_v[idx] = False if m == 0: return keep_v, np.ones(0, dtype=bool) a = edges_int[:, 0].clip(0, n - 1) b = edges_int[:, 1].clip(0, n - 1) keep_e = valid_e & keep_v[a] & keep_v[b] return keep_v, keep_e def build_gt_verts( sample: dict, bbox_center: np.ndarray, bbox_scale: float, k_verts: int = 64, ) -> np.ndarray: """Build padded GT vertex tensor of shape (K, 4) in normalised coords. Column 3 is the validity flag: +1 for real vertex, -1 for padding. No yaw rotation is applied here — the dataloader applies a random rotation to both `scene_xyz` and `verts_gt[:, :3]` per epoch. """ verts = np.array(sample["wf_vertices"], dtype=np.float32) n_real = min(len(verts), k_verts) assert n_real > 0, "scene has zero GT vertices" out = np.zeros((k_verts, 4), dtype=np.float32) verts_norm = (verts[:n_real] - bbox_center) / bbox_scale out[:n_real, :3] = verts_norm.astype(np.float32) out[:n_real, 3] = 1.0 out[n_real:, 3] = -1.0 return out