"""S23DR 2026 — submission script (trained WireframeDiffusion). HF Competitions entry point: the test runner executes ``python script.py`` in the submission repo's working directory. We therefore: * Load ``params.json`` (mirrors ``usm3d/empty_submission_2026``). * Resolve the dataset (test server: ``/tmp/data``; local: snapshot_download). * Stream samples from the webdataset, **load the diffusion model exactly once** on CUDA, run inference, and write ``submission.json``. Online preprocessing -------------------- Uses V10 online preprocessing by default — ``build_scene_input`` produces the V10 feature point cloud (COLMAP+camera tokens + depth-unprojected tokens, per-point RGB, soft top-2 gestalt). Matches the cache layout the model was trained on. Efficiency notes ---------------- * Model + checkpoint are loaded once at module-main, not per worker — a process-pool would otherwise pay full GPU-init cost per sample. * CPU preprocessing (COLMAP parse, depth back-projection, gestalt rasterise) runs in a small ``ThreadPoolExecutor`` so the next sample is ready by the time the GPU finishes the current one. Threads — not processes — keep the model on a single device while NumPy / Pillow / pycolmap drop the GIL. * On any per-sample failure we fall back to the empty solution; the run never aborts mid-dataset. Specifying the checkpoint ------------------------- Resolution order (first hit wins): 1. ``--ckpt PATH`` (CLI, for local testing) 2. ``S23DR_CKPT`` env var (set by SLURM / docker entrypoint) 3. ``params['ckpt']`` (extra field in params.json) 4. ``./test_checkpoint.pth`` (default — what you'd commit to the submission repo) The script auto-resolves model architecture (small vs. large) from the ``args`` dict embedded in the checkpoint, so the same script handles both. Local sanity check ------------------ Before pushing a submission repo, run:: python script.py --sanity This loads the public ``usm3d/hoho22k_2026_trainval`` dataset (which has GT wireframes), runs the trained model on a few validation samples, and prints per-sample + mean HSS / F1 / IoU. It writes no submission.json — its only job is to verify that the checkpoint loads, the V10 online preprocessing path works end-to-end, and the model produces sane output before you submit. """ from __future__ import annotations import argparse import io import json import os import queue import sys import time from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple import numpy as np import torch try: from datasets import load_dataset except ModuleNotFoundError: load_dataset = None from tqdm import tqdm # --- our project imports ---------------------------------------------------- from hoho2025.example_solutions import read_colmap_rec, _cam_matrix_from_image from data.preprocess import ( build_scene_input, _image_to_rgb_array, gestalt_img_to_ids, _camera_for_image, _resolve_colmap_image, GESTALT_CLASSES, ) from data.dataset import _seed_from_order_id from data.stage2_build import build_stage2_scene from data.stage2_dataset import stage2_row_to_sample, _pad_init_verts from data.fps import fps_subsample_stage2_batch from models.scene_encoder import SceneEncoder from models.denoiser import VertexDenoiser from models.diffusion import WireframeDiffusion from models.stage2 import Stage2Diffusion # --------------------------------------------------------------------------- # Empty / safe fallback # --------------------------------------------------------------------------- def empty_solution() -> Tuple[np.ndarray, List[Tuple[int, int]]]: """Minimal valid solution: 2 coincident vertices + 1 edge.""" return np.zeros((2, 3)), [(0, 1)] # --------------------------------------------------------------------------- # Model loading (called once) # --------------------------------------------------------------------------- def resolve_ckpt_path(cli_arg: Optional[str], params: Dict[str, Any]) -> Path: for candidate, src in [ (cli_arg, "--ckpt"), (os.environ.get("S23DR_CKPT"), "$S23DR_CKPT"), (params.get("ckpt"), "params.json[ckpt]"), ("./test_checkpoint.pth", "default ./test_checkpoint.pth"), ]: if candidate: print(f"[ckpt] using {candidate} (from {src})", flush=True) return Path(candidate) raise SystemExit("[ckpt] no checkpoint specified") def resolve_stage2_ckpt_path(cli_arg: Optional[str], params: Dict[str, Any]) -> Path: """Stage-2 checkpoint resolution. The two-stage architecture is mandatory; if no stage-2 ckpt is configured the script aborts. Resolution mirrors ``resolve_ckpt_path``: CLI → env → params.json → committed default.""" for candidate, src in [ (cli_arg, "--stage2_ckpt"), (os.environ.get("S23DR_STAGE2_CKPT"), "$S23DR_STAGE2_CKPT"), (params.get("stage2_ckpt"), "params.json[stage2_ckpt]"), ("./test_checkpoint_stage2.pth", "default ./test_checkpoint_stage2.pth"), ]: if candidate: print(f"[stage2/ckpt] using {candidate} (from {src})", flush=True) return Path(candidate) raise SystemExit("[stage2/ckpt] no stage-2 checkpoint specified — " "two-stage pipeline is mandatory.") def load_model(ckpt_path: Path, device: torch.device): """Return (model, ckpt_args) — model is .eval() on `device`. Architecture hyperparameters (d_model, n_heads, n_layers, d_ff, n_pool, encoder layer counts, ...) are read from the ``args`` dict embedded in the checkpoint. Same loader handles small and large configs transparently. """ print(f"[model] loading {ckpt_path} on {device}", flush=True) t0 = time.perf_counter() ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) ckpt_args = ckpt.get("args", {}) or {} # RGB branch is present in v10+ caches and was used when training models # on those shards. Older sem_v7 checkpoints disable the RGB MLP. cache_version = str(ckpt_args.get("cache_version", "sem_v10")).lower() use_rgb = cache_version not in {"v7", "sem_v7"} enc = SceneEncoder( d_model=ckpt_args.get("d_model", 256), n_heads=ckpt_args.get("n_heads", 8), d_ff=ckpt_args.get("d_ff", 1024), n_full_layers=ckpt_args.get("n_encoder_full_layers", 2), n_pool=ckpt_args.get("n_pool", 1024), n_pool_cross_layers=ckpt_args.get("n_encoder_pool_cross_layers", 1), n_pool_layers=ckpt_args.get("n_encoder_pool_layers", 2), use_rgb=use_rgb, ) den = VertexDenoiser( d_model=ckpt_args.get("d_model", 256), n_heads=ckpt_args.get("n_heads", 8), n_layers=ckpt_args.get("n_layers", 8), d_ff=ckpt_args.get("d_ff", 1024), k_verts=ckpt_args.get("k_verts", 64), ) model = WireframeDiffusion( enc, den, noise_sigma_xyz=ckpt_args.get("noise_sigma_xyz", 1.0), init_from_scene=ckpt_args.get("init_from_scene", False), scene_init_jitter=ckpt_args.get("scene_init_jitter", 0.05), ).to(device) state = ckpt.get("model") or ckpt.get("state_dict") if state is None: raise RuntimeError(f"checkpoint {ckpt_path} has no 'model' or 'state_dict' key") # Strip DDP prefix if present. if any(k.startswith("module.") for k in state): state = {k[len("module."):]: v for k, v in state.items()} missing, unexpected = model.load_state_dict(state, strict=False) if missing or unexpected: print(f"[model] state_dict mismatch: missing={missing} unexpected={unexpected}", flush=True) model.eval() print(f"[model] loaded in {time.perf_counter() - t0:.1f}s " f"d_model={ckpt_args.get('d_model')} n_layers={ckpt_args.get('n_layers')} " f"use_rgb={use_rgb} step={ckpt.get('step')} best_hss={ckpt.get('best_hss')}", flush=True) return model, ckpt_args def load_stage2_model(ckpt_path: Path, device: torch.device): """Build a Stage2Diffusion model from its checkpoint. Architecture is read from the embedded ``args`` dict so the same loader handles small / large variants. Returns (model, ckpt_args). """ print(f"[stage2/model] loading {ckpt_path} on {device}", flush=True) t0 = time.perf_counter() ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) a = ckpt.get("args", {}) or {} if not isinstance(a, dict): a = vars(a) enc = SceneEncoder( d_model=a.get("d_model", 256), n_heads=a.get("n_heads", 8), d_ff=a.get("d_ff", 1024), n_full_layers=a.get("n_encoder_full_layers", 2), n_pool=a.get("n_pool", 1024), n_pool_cross_layers=a.get("n_encoder_pool_cross_layers", 1), n_pool_layers=a.get("n_encoder_pool_layers", 2), use_rgb=True, # stage-2 cache always carries RGB ) den = VertexDenoiser( d_model=a.get("d_model", 256), n_heads=a.get("n_heads", 8), n_layers=a.get("n_layers", 8), d_ff=a.get("d_ff", 1024), k_verts=a.get("k_verts", 64), ) model = Stage2Diffusion( enc, den, noise_sigma_xyz=a.get("noise_sigma_xyz", 0.5), init_from_scene=a.get("init_from_scene", True), scene_init_jitter=a.get("scene_init_jitter", 0.05), ).to(device) state = ckpt.get("model") or ckpt.get("state_dict") if state is None: raise RuntimeError(f"checkpoint {ckpt_path} has no 'model' or 'state_dict' key") if any(k.startswith("module.") for k in state): state = {k[len("module."):]: v for k, v in state.items()} missing, unexpected = model.load_state_dict(state, strict=False) if missing or unexpected: print(f"[stage2/model] state_dict mismatch: missing={missing} unexpected={unexpected}", flush=True) model.eval() print(f"[stage2/model] loaded in {time.perf_counter() - t0:.1f}s " f"d_model={a.get('d_model')} n_layers={a.get('n_layers')} " f"step={ckpt.get('step')} best_hss={ckpt.get('best_hss')}", flush=True) return model, a # --------------------------------------------------------------------------- # Per-sample preprocessing (CPU, runs in worker threads) # --------------------------------------------------------------------------- def preprocess_sample( sample: Dict[str, Any], n_pts: int, use_depth: bool, voxel_size_m: float = 0.1, keep_raw: bool = False, ) -> Dict[str, Any]: """Heavy CPU work: COLMAP parse + V10 scene input build. No torch / GPU here. Mirrors `data.dataset.process_sample` exactly: same `voxel_size_m` (configs set this to 0.1; build_scene_input's bare default is 0.0 which would skip voxel downsampling and produce a different point set) and same per-order_id deterministic RNG seed (so the same points are sampled as were cached at training time). With ``keep_raw=True`` the return value is ``{"scene": ..., "raw_sample": ..., "colmap_rec": ...}`` so the stage-2 builder can reuse the already- decoded COLMAP record instead of parsing the zip twice. """ colmap_rec = read_colmap_rec(sample["colmap"]) order_id = sample.get("order_id", "") rng = np.random.default_rng(_seed_from_order_id(order_id)) scene = build_scene_input( {**sample, "colmap": colmap_rec}, n_pts=n_pts, use_depth=use_depth, voxel_size_m=voxel_size_m, rng=rng, return_cache=keep_raw, ) if not keep_raw: return scene cache = scene.pop("_cache", None) return { "scene": scene, "raw_sample": sample, "colmap_rec": colmap_rec, "cache": cache, } def scene_to_batch(scene: Dict[str, np.ndarray], device: torch.device) -> Dict[str, torch.Tensor]: batch = { "scene_xyz": torch.from_numpy(scene["scene_xyz"]).unsqueeze(0).to(device), "scene_type_ids": torch.from_numpy(scene["scene_type_ids"]).unsqueeze(0).to(device), "scene_gestalt_ids": torch.from_numpy(scene["scene_gestalt_ids"]).unsqueeze(0).to(device), "scene_ade_ids": torch.from_numpy(scene["scene_ade_ids"]).unsqueeze(0).to(device), "bbox_center": torch.from_numpy(scene["bbox_center"]).unsqueeze(0).to(device), "bbox_scale": torch.tensor([scene["bbox_scale"]], device=device), } # Optional V10 fields — included when produced by build_scene_input. for key in ("scene_gestalt_id2", "scene_gestalt_w1", "scene_geom_conf", "scene_sem_conf", "scene_rgb", "bbox_R"): if key in scene: batch[key] = torch.from_numpy(scene[key]).unsqueeze(0).to(device) return batch # --------------------------------------------------------------------------- # Stage-2 inference helpers # --------------------------------------------------------------------------- class Stage2Config: """Frozen knobs for the stage-2 refinement path. Built once in `main`.""" __slots__ = ("model", "n_pts", "k_verts", "n_sample_steps", "validity_thresh", "above_m", "below_m", "side_m", "n_col_oversample", "n_dep_oversample", "n_depth_per_image_cap", "sampling", "density_voxel_size_m", "density_kernel_radius", "density_kernel_axis", "density_response_power", "density_planarity_suppression", "density_planarity_radius", "density_planarity_min_points", "density_min_per_voxel", "ensemble_n", "ensemble_merge_m", "ensemble_mode", "ensemble_strategy", "ensemble_refine_positions", "ensemble_refine_agg", "ensemble_top_k", "ensemble_edge_vote", "ensemble_edge_vote_frac", "ensemble_vertex_vote_frac", "ensemble_stage2_last_k", "ensemble_selector", "ensemble_ranker_path", "ensemble_ranker", "ensemble_fixed_candidate", "ensemble_topk_fuse", "fps_max_exact_iters") def __init__(self, model, n_pts: int, k_verts: int, n_sample_steps: int, validity_thresh: float, above_m: float, below_m: float, side_m: float, n_col_oversample: int, n_dep_oversample: int, n_depth_per_image_cap: int, sampling: str = "random", density_voxel_size_m: float = 0.25, density_kernel_radius: int = 3, density_kernel_axis: str = "cube", density_response_power: float = 1.0, density_planarity_suppression: float = 1.0, density_planarity_radius: int = 1, density_planarity_min_points: int = 12, density_min_per_voxel: int = 0, ensemble_n: int = 1, ensemble_merge_m: float = 0.5, ensemble_mode: str = "medoid", ensemble_strategy: str = "union_hull", ensemble_refine_positions: bool = True, ensemble_refine_agg: str = "mean", ensemble_top_k: int = 0, ensemble_edge_vote: bool = True, ensemble_edge_vote_frac: float = 0.5, ensemble_vertex_vote_frac: float = 0.0, ensemble_stage2_last_k: int = 1, ensemble_selector: str = "legacy", ensemble_ranker_path: Optional[str] = None, ensemble_ranker: Optional[Dict[str, np.ndarray]] = None, ensemble_fixed_candidate: str = "confidence_mean_add", ensemble_topk_fuse: int = 1, fps_max_exact_iters: Optional[int] = None): self.model = model self.n_pts = int(n_pts) self.k_verts = int(k_verts) self.n_sample_steps = int(n_sample_steps) self.validity_thresh = float(validity_thresh) self.above_m = float(above_m); self.below_m = float(below_m); self.side_m = float(side_m) self.n_col_oversample = int(n_col_oversample) self.n_dep_oversample = int(n_dep_oversample) self.n_depth_per_image_cap = int(n_depth_per_image_cap) if sampling not in ("random", "fps", "voxel"): raise ValueError(f"sampling must be 'random', 'fps', or 'voxel', got {sampling!r}") self.sampling = sampling self.density_voxel_size_m = float(density_voxel_size_m) self.density_kernel_radius = int(density_kernel_radius) self.density_kernel_axis = str(density_kernel_axis) self.density_response_power = float(density_response_power) self.density_planarity_suppression = float(density_planarity_suppression) self.density_planarity_radius = int(density_planarity_radius) self.density_planarity_min_points = int(density_planarity_min_points) self.density_min_per_voxel = int(density_min_per_voxel) self.ensemble_n = max(1, int(ensemble_n)) self.ensemble_merge_m = float(ensemble_merge_m) if ensemble_mode not in ("medoid", "consensus", "confidence"): raise ValueError(f"ensemble_mode must be 'medoid', 'consensus', or 'confidence', got {ensemble_mode!r}") self.ensemble_mode = str(ensemble_mode) if ensemble_strategy not in ("union_hull", "per_seed"): raise ValueError(f"ensemble_strategy must be 'union_hull' or 'per_seed', got {ensemble_strategy!r}") self.ensemble_strategy = str(ensemble_strategy) self.ensemble_refine_positions = bool(ensemble_refine_positions) if ensemble_refine_agg not in ("median", "mean", "wmean"): raise ValueError( f"ensemble_refine_agg must be 'median', 'mean', or 'wmean', " f"got {ensemble_refine_agg!r}") self.ensemble_refine_agg = str(ensemble_refine_agg) self.ensemble_top_k = max(0, int(ensemble_top_k)) self.ensemble_edge_vote = bool(ensemble_edge_vote) if not (0.0 < float(ensemble_edge_vote_frac) <= 1.0): raise ValueError( f"ensemble_edge_vote_frac must be in (0, 1], got {ensemble_edge_vote_frac!r}") self.ensemble_edge_vote_frac = float(ensemble_edge_vote_frac) if not (0.0 <= float(ensemble_vertex_vote_frac) <= 1.0): raise ValueError( f"ensemble_vertex_vote_frac must be in [0, 1], " f"got {ensemble_vertex_vote_frac!r}") self.ensemble_vertex_vote_frac = float(ensemble_vertex_vote_frac) self.ensemble_stage2_last_k = max(1, int(ensemble_stage2_last_k)) if ensemble_selector not in ("legacy", "ranker", "fixed"): raise ValueError( f"ensemble_selector must be 'legacy', 'ranker', or 'fixed', got {ensemble_selector!r}") self.ensemble_selector = str(ensemble_selector) self.ensemble_ranker_path = ensemble_ranker_path self.ensemble_ranker = ensemble_ranker self.ensemble_fixed_candidate = str(ensemble_fixed_candidate) self.ensemble_topk_fuse = max(1, int(ensemble_topk_fuse)) # None => full exact FPS. Set to the value the model trained with so # the per-provenance anchor count matches; otherwise the encoder sees # a markedly different point distribution at inference (all exact # anchors vs. exact prefix + deterministic random fill). self.fps_max_exact_iters = (None if fps_max_exact_iters is None else int(fps_max_exact_iters)) def _stage2_sample_to_batch(sample: Dict[str, Any], device: torch.device) -> Dict[str, torch.Tensor]: """Same shape as ``scene_to_batch`` but for the stage-2 sample dict produced by ``stage2_row_to_sample`` (includes ``init_verts`` + valid mask).""" keys = ( "scene_xyz", "scene_type_ids", "scene_gestalt_ids", "scene_gestalt_id2", "scene_gestalt_w1", "scene_ade_ids", "scene_geom_conf", "scene_sem_conf", "scene_rgb", "init_verts", "init_verts_valid", "bbox_center", "bbox_scale", "bbox_R", # FPS-mode extras (only present when pre_subsample=False): "scene_valid_mask", "init_verts_world", "verts_gt_world", ) return {k: sample[k].unsqueeze(0).to(device) for k in keys if k in sample} def _replicate_batch(batch: Dict[str, torch.Tensor], n: int) -> Dict[str, torch.Tensor]: """Return a shallow-copied batch with every leading-dim-1 tensor repeated to leading dim `n`. Non-tensor / non-(1, ...) entries pass through. Used to convert a B=1 single-scene batch into a B=N replicated batch so the diffusion sampler produces N independent trajectories sharing one scene.""" if n <= 1: return batch out: Dict[str, torch.Tensor] = {} for k, v in batch.items(): if isinstance(v, torch.Tensor) and v.dim() >= 1 and v.shape[0] == 1: out[k] = v.repeat(n, *([1] * (v.dim() - 1))) else: out[k] = v return out def _stage1_seed_ensemble( model: WireframeDiffusion, batch1: Dict[str, torch.Tensor], n_steps: int, validity_thresh: float, n_seeds: int, ) -> List[Tuple[np.ndarray, np.ndarray, np.ndarray]]: """Run stage-1 sampling `n_seeds` times with independent random init x0. Mirrors the body of ``WireframeDiffusion.sample`` so we can extract the final-step validity logits per slot (needed for confidence-mode ensemble selection); ``model.sample`` only returns a boolean mask and drops the raw logits. Returns a list of ``(verts_world (M, 3) float32, edges (E, 2) int64, valid_logits (M,) float32)`` tuples, one per seed. ``valid_logits`` contains the raw logit value for each kept vertex (post-threshold) so callers can compute mean/max self-confidence.""" n_seeds = max(1, int(n_seeds)) batch_n = _replicate_batch(batch1, n_seeds) if n_seeds > 1 else batch1 device = batch_n["scene_xyz"].device # Inline equivalent of model.sample(), but keeping last_logit available. scene_feats, scene_xyz, query_xyz = model._encode_scene_with_query(batch_n) K = model.denoiser.k_verts B = scene_feats.shape[0] dt = 1.0 / max(1, n_steps) x = model._init_x0(batch_n, K, device, B, query_xyz=query_xyz) last_logit = torch.zeros(B, K, device=device) last_edge_logit = torch.zeros(B, K, K, device=device) for i in range(n_steps): t = torch.full((B,), i * dt, device=device) v, last_logit, last_edge_logit = model.denoiser( x, t, scene_feats, scene_xyz, ) v = v.float() last_logit = last_logit.float() last_edge_logit = last_edge_logit.float() x = x + dt * v out: List[Tuple[np.ndarray, np.ndarray, np.ndarray]] = [] bbox_R = batch_n.get("bbox_R") for i in range(B): x_i = torch.nan_to_num(x[i].float(), nan=0.0, posinf=model.xyz_clip, neginf=-model.xyz_clip ).clamp(-model.xyz_clip, model.xyz_clip) logit_i = last_logit[i] idx = torch.nonzero(logit_i > validity_thresh, as_tuple=False).flatten() x_valid = x_i.index_select(0, idx) center = batch_n["bbox_center"][i].to(device=x_valid.device, dtype=torch.float32) scale = batch_n["bbox_scale"][i].to(device=x_valid.device, dtype=torch.float32) if isinstance(bbox_R, torch.Tensor): R = bbox_R[i].to(device=x_valid.device, dtype=torch.float32) verts_world = (x_valid * scale) @ R + center else: verts_world = x_valid * scale + center verts_np = verts_world.float().cpu().numpy().reshape(-1, 3).astype(np.float32) logits_kept = logit_i.index_select(0, idx).detach().cpu().numpy().astype(np.float32) n_valid = int(idx.numel()) edges_arr = np.zeros((0, 2), dtype=np.int64) if n_valid >= 2: sub = last_edge_logit[i].float().index_select(0, idx).index_select(1, idx) tri = torch.triu(torch.ones(n_valid, n_valid, device=sub.device, dtype=torch.bool), diagonal=1) pairs = torch.nonzero((sub > 0.0) & tri, as_tuple=False) if pairs.numel() > 0: edges_arr = pairs.detach().cpu().numpy().astype(np.int64) out.append((verts_np, edges_arr, logits_kept)) return out def _fuse_wireframes( runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], n_total: int, merge_tau_m: float, vertex_vote_frac: float = 0.5, edge_vote_frac: float = 0.5, aggregator: str = "mean", ) -> Tuple[np.ndarray, List[Tuple[int, int]]]: """Vertex-merge consensus fusion. 1. Greedy single-link cluster all run vertices in world coords (threshold ``merge_tau_m`` metres). Each cluster centroid is the running mean of its members (or a logit-weighted mean when ``aggregator="wmean"``). 2. Keep clusters supported by at least ``ceil(n_total * vertex_vote_frac)`` distinct runs. 3. For each pair of kept clusters, vote an edge if it appears in at least ``ceil(n_total * edge_vote_frac)`` runs. Falls back to the single run with the most predicted vertices when no cluster meets the support bar — we never collapse to empty when individual runs were non-empty. ``runs`` entries are ``(verts, edges, logits)``; ``logits`` is only consulted under ``aggregator="wmean"``.""" if aggregator not in ("mean", "median", "wmean"): raise ValueError( f"aggregator must be 'mean', 'median', or 'wmean', got {aggregator!r}") valid: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]] = [] for r in runs: v_raw = r[0] if v_raw is None or len(v_raw) == 0: continue v = np.asarray(v_raw, dtype=np.float64).reshape(-1, 3) e = list(r[1]) if len(r) > 1 and r[1] is not None else [] lg_raw = r[2] if len(r) > 2 else None lg = (np.asarray(lg_raw, dtype=np.float64).reshape(-1) if lg_raw is not None else np.zeros(v.shape[0], dtype=np.float64)) if lg.shape[0] != v.shape[0]: lg = np.zeros(v.shape[0], dtype=np.float64) valid.append((v, e, lg)) if not valid: return np.zeros((0, 3), dtype=np.float64), [] v_support = max(1, int(np.ceil(n_total * float(vertex_vote_frac)))) e_support = max(1, int(np.ceil(n_total * float(edge_vote_frac)))) # Per-cluster bookkeeping: member positions, member logits, and the set # of run indices contributing. Centroid is computed once at the end so # 'median' / 'wmean' can use all members (running-mean only worked for # the plain 'mean' aggregator). cluster_members: List[List[np.ndarray]] = [] cluster_member_lg: List[List[float]] = [] cluster_centroids: List[np.ndarray] = [] # running mean — used for matching cluster_count: List[int] = [] cluster_runs: List[set] = [] vertex_to_cluster: List[List[int]] = [] for r, (verts, _edges, lg) in enumerate(valid): local_map: List[int] = [] for vi in range(verts.shape[0]): v = verts[vi] best_c, best_d = -1, float("inf") for ci, cc in enumerate(cluster_centroids): d = float(np.linalg.norm(v - cc)) if d < best_d: best_d, best_c = d, ci if best_c >= 0 and best_d <= merge_tau_m: n = cluster_count[best_c] cluster_centroids[best_c] = (cluster_centroids[best_c] * n + v) / (n + 1) cluster_count[best_c] = n + 1 cluster_runs[best_c].add(r) cluster_members[best_c].append(v) cluster_member_lg[best_c].append(float(lg[vi])) local_map.append(best_c) else: cluster_centroids.append(v.copy()) cluster_count.append(1) cluster_runs.append({r}) cluster_members.append([v]) cluster_member_lg.append([float(lg[vi])]) local_map.append(len(cluster_centroids) - 1) vertex_to_cluster.append(local_map) kept = [ci for ci, rs in enumerate(cluster_runs) if len(rs) >= v_support] if not kept: best = max(range(len(valid)), key=lambda r: valid[r][0].shape[0]) return valid[best][0].astype(np.float64), list(valid[best][1]) # Final centroids per kept cluster, per aggregator. fused_verts = np.empty((len(kept), 3), dtype=np.float64) for new_i, ci in enumerate(kept): stack = np.stack(cluster_members[ci], axis=0) if aggregator == "median": fused_verts[new_i] = np.median(stack, axis=0) elif aggregator == "wmean": lg = np.asarray(cluster_member_lg[ci], dtype=np.float64) lg = lg - lg.max() w = np.exp(lg) ws = float(w.sum()) if ws <= 1e-12 or not np.isfinite(ws): fused_verts[new_i] = stack.mean(axis=0) else: fused_verts[new_i] = (stack * (w / ws)[:, None]).sum(axis=0) else: # "mean" fused_verts[new_i] = stack.mean(axis=0) old_to_new = {old: new for new, old in enumerate(kept)} n_kept = len(kept) edge_votes = np.zeros((n_kept, n_kept), dtype=np.int32) for r, (_verts, edges, _lg) in enumerate(valid): seen_pairs: set = set() for (a, b) in edges: if a < 0 or b < 0 or a == b: continue if a >= len(vertex_to_cluster[r]) or b >= len(vertex_to_cluster[r]): continue ca = vertex_to_cluster[r][a] cb = vertex_to_cluster[r][b] if ca == cb or ca not in old_to_new or cb not in old_to_new: continue i, j = old_to_new[ca], old_to_new[cb] if i > j: i, j = j, i if (i, j) in seen_pairs: continue seen_pairs.add((i, j)) edge_votes[i, j] += 1 ii, jj = np.where(edge_votes >= e_support) fused_edges = [(int(a), int(b)) for a, b in zip(ii.tolist(), jj.tolist())] return fused_verts, fused_edges def _wireframe_pair_distance( va: np.ndarray, ea: List[Tuple[int, int]], vb: np.ndarray, eb: List[Tuple[int, int]], match_tau_m: float = 0.4, ) -> float: """Symmetric distance between two wireframes for medoid selection. Uses Hungarian (optimal 1-1) vertex assignment, which is naturally symmetric and aligned with how the HSS metric matches predictions to GT. Components: * **Vertex term** (metres, capped at ``match_tau_m``): solve the rectangular linear-sum-assignment on the pairwise distance matrix capped at ``match_tau_m``. Matched pair cost = capped distance; unmatched leftover vertices (size mismatch) cost ``match_tau_m`` each. Mean over ``max(|Va|, |Vb|)`` slots so the term lives in ``[0, match_tau_m]``. * **Edge term** (Jaccard, in ``[0, 1]``): the Hungarian assignment gives a 1-1 vertex map; use it to relabel B's edges into A's index space (only for matches actually within ``match_tau_m``) and compute ``1 - IoU`` over the two edge sets. Returns a single scalar; lower = more similar. Degenerate empty inputs return ``+inf`` so they cannot be picked as the medoid. """ from scipy.optimize import linear_sum_assignment va = np.asarray(va, dtype=np.float64).reshape(-1, 3) vb = np.asarray(vb, dtype=np.float64).reshape(-1, 3) Na, Nb = va.shape[0], vb.shape[0] if Na == 0 or Nb == 0: return float("inf") D = np.linalg.norm(va[:, None, :] - vb[None, :, :], axis=-1) D_capped = np.minimum(D, match_tau_m) row_ind, col_ind = linear_sum_assignment(D_capped) matched_cost = float(D_capped[row_ind, col_ind].sum()) n_unmatched = abs(Na - Nb) # leftover on the larger side vertex_dist = (matched_cost + n_unmatched * match_tau_m) / max(Na, Nb) # Map B → A via the Hungarian assignment, but only when the match is # actually within tau (capped matches at exactly tau are unreliable). b_to_a = -np.ones(Nb, dtype=np.int64) for r, c in zip(row_ind, col_ind): if D[r, c] <= match_tau_m: b_to_a[c] = r ea_norm = set((min(int(i), int(j)), max(int(i), int(j))) for i, j in ea if i != j) eb_mapped = set() for (i, j) in eb: if i == j or i < 0 or j < 0 or i >= Nb or j >= Nb: continue ai, aj = int(b_to_a[i]), int(b_to_a[j]) if ai < 0 or aj < 0 or ai == aj: continue eb_mapped.add((min(ai, aj), max(ai, aj))) if not ea_norm and not eb_mapped: edge_dist = 0.0 elif not ea_norm or not eb_mapped: edge_dist = 1.0 else: inter = len(ea_norm & eb_mapped) union = len(ea_norm | eb_mapped) edge_dist = 1.0 - (inter / union) return vertex_dist + edge_dist def _medoid_rank( runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], match_tau_m: float = 0.4, ) -> List[int]: """Return indices into ``runs`` sorted by centrality (lowest pairwise Hungarian distance first → most-central run is index ``[0]``). Empty / edgeless runs are excluded from the pairwise computation; if none qualify the singleton list ``[longest_run_index]`` is returned. Always returns at least one valid index when ``runs`` is non-empty.""" if not runs: return [] candidates_idx = [r for r, (v, e, *_) in enumerate(runs) if len(v) > 0 and len(e) > 0] if not candidates_idx: return [max(range(len(runs)), key=lambda r: (len(runs[r][0]) if runs[r][0] is not None else -1, len(runs[r][1]) if runs[r][1] is not None else -1))] if len(candidates_idx) == 1: return [candidates_idx[0]] n = len(candidates_idx) verts = [np.asarray(runs[idx][0], dtype=np.float64).reshape(-1, 3) for idx in candidates_idx] edges = [list(runs[idx][1]) for idx in candidates_idx] dist = np.zeros((n, n), dtype=np.float64) for i in range(n): for j in range(i + 1, n): d = _wireframe_pair_distance( verts[i], edges[i], verts[j], edges[j], match_tau_m=match_tau_m, ) dist[i, j] = d dist[j, i] = d centrality = dist.sum(axis=1) order = np.argsort(centrality, kind="stable") return [candidates_idx[int(k)] for k in order] def _medoid_select_idx( runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], match_tau_m: float = 0.4, ) -> int: """Return the index into ``runs`` of the medoid — the run whose wireframe is most similar to the others by Hungarian-matched distance. Thin wrapper around ``_medoid_rank`` that returns just the most-central index.""" ranking = _medoid_rank(runs, match_tau_m=match_tau_m) return ranking[0] if ranking else 0 def _confidence_select_idx( runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], ) -> int: """Return the index into ``runs`` of the highest mean-validity-logit run. Skips empty / all-NaN runs in the scoring; falls back to the longest run when no run has meaningful confidence.""" if not runs: return 0 scored: List[Tuple[float, int]] = [] for i, (v, _e, logits) in enumerate(runs): if v is None or len(v) == 0 or logits is None or len(logits) == 0: continue arr = np.asarray(logits, dtype=np.float64) if not np.all(np.isfinite(arr)): arr = arr[np.isfinite(arr)] if arr.size == 0: continue scored.append((float(arr.mean()), i)) if not scored: return max(range(len(runs)), key=lambda r: (len(runs[r][0]) if runs[r][0] is not None else -1, len(runs[r][1]) if runs[r][1] is not None else -1)) return max(scored, key=lambda t: t[0])[1] def _refine_positions( runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], picked_idx: int, match_tau_m: float, aggregator: str = "mean", member_indices: Optional[Sequence[int]] = None, vertex_vote_frac: float = 0.0, ) -> Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]: """Position refinement on top of a single picked run. Keep the picked run's vertices and edges (topology unchanged). For each vertex ``v`` of the picked run, find its Hungarian-matched correspondent in every other run (only matches within ``match_tau_m`` count), then replace ``v`` with the per-axis ``aggregator`` of ``{v}`` ∪ correspondents. ``aggregator`` ∈ {"mean", "median", "wmean"}: * ``"mean"`` (default) — simple average; lower variance under iid noise but sensitive to outliers. * ``"median"`` — per-axis median. Robust to a single bad correspondent (e.g., an off-by-one Hungarian match at the τ boundary). * ``"wmean"`` — softmax-of-validity-logit weighted mean. Rationale: vertex prediction is the dominant noise source (F1 hits HSS harder than IoU). Selectors (medoid / confidence) pick a *vertex set* that's already good; this step denoises each vertex's *position* by aggregating with consensus detections from the other N−1 runs. Topology (edges) is untouched so we never lose recall on edges. Returns ``(verts_refined, edges_picked)``. Falls back to the picked run's verts/edges verbatim when no refinement signal is available.""" from scipy.optimize import linear_sum_assignment if aggregator not in ("median", "mean", "wmean"): raise ValueError( f"aggregator must be 'median', 'mean', or 'wmean', got {aggregator!r}") if not runs: return np.zeros((0, 3), dtype=np.float64), [], np.zeros((0,), dtype=np.int64) picked_v_raw, picked_e_raw, picked_logits = runs[picked_idx] picked_v = np.asarray(picked_v_raw, dtype=np.float64).reshape(-1, 3) picked_e = list(picked_e_raw) if picked_v.shape[0] == 0: return picked_v, picked_e, np.zeros((0,), dtype=np.int64) M = picked_v.shape[0] picked_lg = (np.asarray(picked_logits, dtype=np.float64).reshape(-1) if picked_logits is not None else np.zeros(M, dtype=np.float64)) if picked_lg.shape[0] != M: picked_lg = np.zeros(M, dtype=np.float64) # Per-picked-vertex correspondent positions and per-correspondent logits # (used for 'wmean' aggregation). Picked run is always the first entry. correspondents: List[List[np.ndarray]] = [[picked_v[i]] for i in range(M)] corr_logits: List[List[float]] = [[float(picked_lg[i])] for i in range(M)] if member_indices is None: contributor_iter: List[int] = list(range(len(runs))) else: contributor_iter = list(member_indices) if picked_idx not in contributor_iter: contributor_iter.append(picked_idx) # Voter set includes picked (self-vote on each vertex). M_voters = max(1, len(set(contributor_iter))) for j in contributor_iter: if j == picked_idx or j < 0 or j >= len(runs): continue vj_raw, _ej, lj_raw = runs[j] vj = np.asarray(vj_raw, dtype=np.float64).reshape(-1, 3) if vj.shape[0] == 0: continue lj = (np.asarray(lj_raw, dtype=np.float64).reshape(-1) if lj_raw is not None else np.zeros(vj.shape[0], dtype=np.float64)) if lj.shape[0] != vj.shape[0]: lj = np.zeros(vj.shape[0], dtype=np.float64) D = np.linalg.norm(picked_v[:, None, :] - vj[None, :, :], axis=-1) D_capped = np.minimum(D, match_tau_m) row_ind, col_ind = linear_sum_assignment(D_capped) for r, c in zip(row_ind, col_ind): if D[r, c] <= match_tau_m: correspondents[r].append(vj[c]) corr_logits[r].append(float(lj[c])) # Vertex consensus drop. A picked vertex survives iff at least # ceil(M_voters * vertex_vote_frac) voters support it (including the # picked run's self-vote). if vertex_vote_frac > 0.0: threshold = int(np.ceil(M_voters * float(vertex_vote_frac))) keep_mask = np.array( [len(correspondents[i]) >= threshold for i in range(M)], dtype=bool, ) if not keep_mask.any(): # Don't collapse to empty — fall back to keeping all picked verts. keep_mask[:] = True else: keep_mask = np.ones(M, dtype=bool) keep_idx = np.flatnonzero(keep_mask).astype(np.int64) refined = np.empty((keep_idx.size, 3), dtype=np.float64) for new_i, old_i in enumerate(keep_idx): stack = np.stack(correspondents[int(old_i)], axis=0) if aggregator == "median": refined[new_i] = np.median(stack, axis=0) elif aggregator == "mean": refined[new_i] = stack.mean(axis=0) else: # "wmean" lg = np.asarray(corr_logits[int(old_i)], dtype=np.float64) # Softmax over logits for stable positive weights; falls back to # uniform when logits are all-equal / missing. lg = lg - lg.max() w = np.exp(lg) w_sum = float(w.sum()) if w_sum <= 1e-12 or not np.isfinite(w_sum): refined[new_i] = stack.mean(axis=0) else: w = w / w_sum refined[new_i] = (stack * w[:, None]).sum(axis=0) # Remap picked-run edges into the kept-vertex index space; drop edges # whose endpoints didn't survive the vertex vote. old_to_new = -np.ones(M, dtype=np.int64) for new_i, old_i in enumerate(keep_idx): old_to_new[int(old_i)] = new_i edges_remapped: List[Tuple[int, int]] = [] for (a, b) in picked_e: ia, ib = int(a), int(b) if ia < 0 or ib < 0 or ia >= M or ib >= M or ia == ib: continue na, nb = int(old_to_new[ia]), int(old_to_new[ib]) if na < 0 or nb < 0 or na == nb: continue edges_remapped.append((na, nb)) return refined, edges_remapped, keep_idx def _vote_edges( runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], picked_idx: int, match_tau_m: float, vote_frac: float = 0.5, member_indices: Optional[Sequence[int]] = None, keep_idx: Optional[np.ndarray] = None, ) -> List[Tuple[int, int]]: """Majority-vote edge refinement. For each voter run, Hungarian-match its vertices to the picked run's vertices (matches outside ``match_tau_m`` ignored), relabel its edges into picked-vertex index space, and tally votes per ``(i, j)`` pair. An edge survives if at least ``ceil(M * vote_frac)`` voters (out of ``M`` = ``len(member_indices)`` or ``len(runs)``) vote for it. The picked run itself is always a voter — its own edges therefore enter with one self-vote. Returns the filtered edge list in picked-vertex index space. """ from scipy.optimize import linear_sum_assignment if not runs: return [] picked_v = np.asarray(runs[picked_idx][0], dtype=np.float64).reshape(-1, 3) if picked_v.shape[0] == 0: return list(runs[picked_idx][1]) if member_indices is None: voters: List[int] = list(range(len(runs))) else: voters = [int(j) for j in member_indices if 0 <= int(j) < len(runs)] if picked_idx not in voters: voters.append(int(picked_idx)) M = max(1, len(voters)) threshold = int(np.ceil(M * float(vote_frac))) votes: Dict[Tuple[int, int], int] = {} Mp = picked_v.shape[0] for j in voters: vj_raw, ej_raw, _lj = runs[j] vj = np.asarray(vj_raw, dtype=np.float64).reshape(-1, 3) if j == picked_idx: j_to_p = np.arange(Mp, dtype=np.int64) else: if vj.shape[0] == 0: continue D = np.linalg.norm(picked_v[:, None, :] - vj[None, :, :], axis=-1) D_capped = np.minimum(D, match_tau_m) row_ind, col_ind = linear_sum_assignment(D_capped) j_to_p = -np.ones(vj.shape[0], dtype=np.int64) for r, c in zip(row_ind, col_ind): if D[r, c] <= match_tau_m: j_to_p[c] = r for (a, b) in ej_raw: ia, ib = int(a), int(b) if ia == ib or ia < 0 or ib < 0 or ia >= len(j_to_p) or ib >= len(j_to_p): continue pa, pb = int(j_to_p[ia]), int(j_to_p[ib]) if pa < 0 or pb < 0 or pa == pb: continue key = (min(pa, pb), max(pa, pb)) votes[key] = votes.get(key, 0) + 1 kept_old = [(int(i), int(j)) for (i, j), c in votes.items() if c >= threshold] if not kept_old: # Degenerate vote tally (e.g. only the picked run voted on edges with # vote_frac > 1/M). Fall back to picked edges so we never lose all # edges — F1 would zero out HSS. kept_old = [(int(a), int(b)) for (a, b) in runs[picked_idx][1] if int(a) != int(b)] if keep_idx is None: return kept_old # Remap from picked-vertex index space → kept-vertex index space (the # output of _refine_positions when vertex_vote_frac > 0 drops some # picked vertices). Edges referencing dropped vertices are filtered out. old_to_new = -np.ones(Mp, dtype=np.int64) for new_i, old_i in enumerate(np.asarray(keep_idx).reshape(-1)): oi = int(old_i) if 0 <= oi < Mp: old_to_new[oi] = new_i out: List[Tuple[int, int]] = [] for (a, b) in kept_old: na, nb = int(old_to_new[a]), int(old_to_new[b]) if na < 0 or nb < 0 or na == nb: continue out.append((min(na, nb), max(na, nb))) return out # ------------------------------------------------------------------ # 2D gestalt-mask reprojection features # ------------------------------------------------------------------ # Pixels of these gestalt classes mark wireframe edges (lines). _GESTALT_EDGE_NAMES = ( "ridge", "eave", "rake", "hip", "valley", "fascia", "transition_line", "flashing", "step_flashing", "ground_line", "soffit", ) # Pixels of these classes mark wireframe vertices (line endpoints). _GESTALT_VERT_NAMES = ( "apex", "eave_end_point", "flashing_end_point", ) def _ids_from_class_names(names: Tuple[str, ...]) -> np.ndarray: return np.array( sorted(GESTALT_CLASSES.index(n) for n in names if n in GESTALT_CLASSES), dtype=np.int64, ) _EDGE_GESTALT_IDS = _ids_from_class_names(_GESTALT_EDGE_NAMES) _VERT_GESTALT_IDS = _ids_from_class_names(_GESTALT_VERT_NAMES) def _build_reprojection_views( raw_sample: Dict[str, Any], colmap_rec, edge_dilate: int = 2, vert_dilate: int = 4, ) -> List[Dict[str, Any]]: """Pre-compute per-view camera + dilated gestalt masks for reprojection. Returns a list of dicts with keys ``R`` (3,3), ``t`` (3,), ``fx``, ``fy``, ``cx``, ``cy``, ``W``, ``H``, ``edge_mask`` and ``vert_mask`` (both uint8 HxW). The masks are 1 where pixels belong to wireframe edge / vertex gestalt classes (dilated by ``edge_dilate`` / ``vert_dilate`` 3x3 iterations to tolerate small projection error and label boundary noise). """ try: from scipy.ndimage import binary_dilation except Exception: binary_dilation = None # type: ignore image_ids = raw_sample.get("image_ids", []) or [] gestalt_imgs = raw_sample.get("gestalt", []) or [] if not image_ids or not gestalt_imgs: return [] colmap_by_name = { col_img.name: col_img for col_img in colmap_rec.images.values() } structure = np.ones((3, 3), dtype=bool) out: List[Dict[str, Any]] = [] for i, img_id in enumerate(image_ids): if i >= len(gestalt_imgs) or gestalt_imgs[i] is None: continue col_img = _resolve_colmap_image(colmap_by_name, img_id) if col_img is None: continue cam = _camera_for_image(colmap_rec, col_img) cam_w = int(getattr(cam, "width", 0) or 0) cam_h = int(getattr(cam, "height", 0) or 0) if cam_w <= 0 or cam_h <= 0: continue try: R, t = _cam_matrix_from_image(col_img) R = np.asarray(R, dtype=np.float64).reshape(3, 3) t = np.asarray(t, dtype=np.float64).reshape(3) K = cam.calibration_matrix() except Exception: continue try: gest_np = _image_to_rgb_array(gestalt_imgs[i]) except Exception: continue ids = gestalt_img_to_ids(gest_np) edge_tight = np.isin(ids, _EDGE_GESTALT_IDS) vert_tight = np.isin(ids, _VERT_GESTALT_IDS) edge_mask = edge_tight.copy() vert_mask = vert_tight.copy() if binary_dilation is not None: if edge_dilate > 0 and edge_mask.any(): edge_mask = binary_dilation( edge_mask, structure=structure, iterations=int(edge_dilate)) if vert_dilate > 0 and vert_mask.any(): vert_mask = binary_dilation( vert_mask, structure=structure, iterations=int(vert_dilate)) H, W = ids.shape[:2] out.append({ "R": R, "t": t, "fx": float(K[0, 0]), "fy": float(K[1, 1]), "cx": float(K[0, 2]), "cy": float(K[1, 2]), "cam_w": float(cam_w), "cam_h": float(cam_h), "W": int(W), "H": int(H), "edge_mask": np.ascontiguousarray(edge_mask, dtype=bool), "vert_mask": np.ascontiguousarray(vert_mask, dtype=bool), "edge_mask_tight": np.ascontiguousarray(edge_tight, dtype=bool), "vert_mask_tight": np.ascontiguousarray(vert_tight, dtype=bool), "ids": np.ascontiguousarray(ids, dtype=np.int32), }) return out def _project_to_mask( pts: np.ndarray, view: Dict[str, Any], which: str, ) -> Tuple[np.ndarray, np.ndarray]: """Project (N,3) world points into one view's mask. Returns (in_view, hit) bool arrays of length N. ``in_view`` is True when the point is in front of the camera and lands inside the image; ``hit`` is True when in_view AND the mask is set at the corresponding pixel. """ R = view["R"]; t = view["t"] p_cam = pts @ R.T + t z = p_cam[:, 2] valid = z > 1e-6 u = np.empty_like(z); v = np.empty_like(z) zsafe = np.where(valid, z, 1.0) u_proj = p_cam[:, 0] / zsafe * view["fx"] + view["cx"] v_proj = p_cam[:, 1] / zsafe * view["fy"] + view["cy"] sx = view["W"] / view["cam_w"]; sy = view["H"] / view["cam_h"] ui = np.rint(u_proj * sx).astype(np.int64) vi = np.rint(v_proj * sy).astype(np.int64) in_view = valid & (ui >= 0) & (ui < view["W"]) & (vi >= 0) & (vi < view["H"]) mask = view[which] hit = np.zeros(pts.shape[0], dtype=bool) if in_view.any(): idx = np.where(in_view)[0] hit[idx] = mask[vi[idx], ui[idx]] return in_view, hit def _reprojection_features( verts: np.ndarray, edges: List[Tuple[int, int]], views: List[Dict[str, Any]], n_edge_samples: int = 8, ) -> Tuple[float, float, float, float]: """Return ``(vert_mean, vert_min, edge_mean, edge_min)`` — legacy stub. Mostly kept for backwards-compatibility; new code should call :func:`_reprojection_features_extended` for the richer per-vertex / per-edge breakdown. """ res = _reprojection_features_extended(verts, edges, views, n_edge_samples) return res["vert_mean"], res["vert_min"], res["edge_mean"], res["edge_min"] def _reprojection_features_extended( verts: np.ndarray, edges: List[Tuple[int, int]], views: List[Dict[str, Any]], n_edge_samples: int = 8, ) -> Dict[str, Any]: """Full reprojection scoring with per-vertex / per-edge arrays. Returns a dict with scalar fields: vert_mean, vert_min, vert_worst3_mean, vert_bad_frac, vert_soft_mean, edge_mean, edge_min, edge_worst3_mean, edge_bad_frac, edge_soft_mean, endpoint_corner_mean, endpoint_corner_min, plus arrays: per_v (Nv,), per_e (Ne,) All scalars in [0, 1]. "soft" uses the non-dilated (tight) masks; "bad" thresholds at 0.1; "worst3" averages the 3 lowest scores. """ out = { "vert_mean": 0.0, "vert_min": 0.0, "vert_worst3_mean": 0.0, "vert_bad_frac": 0.0, "vert_soft_mean": 0.0, "edge_mean": 0.0, "edge_min": 0.0, "edge_worst3_mean": 0.0, "edge_bad_frac": 0.0, "edge_soft_mean": 0.0, "endpoint_corner_mean": 0.0, "endpoint_corner_min": 0.0, "per_v": np.zeros((0,), dtype=np.float64), "per_e": np.zeros((0,), dtype=np.float64), } if not views or verts.size == 0: return out v = np.asarray(verts, dtype=np.float64).reshape(-1, 3) Nv = v.shape[0] # --- vertices: dilated + tight masks --- vert_hits = np.zeros(Nv, dtype=np.float64) vert_soft = np.zeros(Nv, dtype=np.float64) vert_views = np.zeros(Nv, dtype=np.float64) for view in views: in_view, hit = _project_to_mask(v, view, "vert_mask") _, hit_soft = _project_to_mask(v, view, "vert_mask_tight") vert_views += in_view.astype(np.float64) vert_hits += hit.astype(np.float64) vert_soft += hit_soft.astype(np.float64) denom = np.maximum(vert_views, 1.0) per_v = np.where(vert_views > 0, vert_hits / denom, 0.0) per_v_soft = np.where(vert_views > 0, vert_soft / denom, 0.0) out["per_v"] = per_v out["vert_mean"] = float(per_v.mean()) out["vert_min"] = float(per_v.min()) k3 = min(3, Nv) out["vert_worst3_mean"] = float(np.sort(per_v)[:k3].mean()) if k3 else 0.0 out["vert_bad_frac"] = float(np.mean(per_v < 0.1)) out["vert_soft_mean"] = float(per_v_soft.mean()) if not edges: return out e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2) Ne = e_arr.shape[0] # --- edges: sample interior, dilated + tight masks --- ts = np.linspace(0.0, 1.0, n_edge_samples + 2)[1:-1] a = v[e_arr[:, 0]]; b = v[e_arr[:, 1]] samples = a[:, None, :] + (b - a)[:, None, :] * ts[None, :, None] sflat = samples.reshape(-1, 3) eh = np.zeros(Ne * n_edge_samples, dtype=np.float64) eh_soft = np.zeros_like(eh) ev = np.zeros_like(eh) for view in views: in_view, hit = _project_to_mask(sflat, view, "edge_mask") _, hit_soft = _project_to_mask(sflat, view, "edge_mask_tight") ev += in_view.astype(np.float64) eh += hit.astype(np.float64) eh_soft += hit_soft.astype(np.float64) edenom = np.maximum(ev, 1.0) per_sample = np.where(ev > 0, eh / edenom, 0.0) per_sample_soft = np.where(ev > 0, eh_soft / edenom, 0.0) per_e = per_sample.reshape(Ne, n_edge_samples).mean(axis=1) per_e_soft = per_sample_soft.reshape(Ne, n_edge_samples).mean(axis=1) out["per_e"] = per_e out["edge_mean"] = float(per_e.mean()) out["edge_min"] = float(per_e.min()) k3e = min(3, Ne) out["edge_worst3_mean"] = float(np.sort(per_e)[:k3e].mean()) if k3e else 0.0 out["edge_bad_frac"] = float(np.mean(per_e < 0.1)) out["edge_soft_mean"] = float(per_e_soft.mean()) # --- endpoint-on-corner: per-edge min(endpoint vertex scores) --- if Nv > 0: e_endpoint_min = np.minimum(per_v[e_arr[:, 0]], per_v[e_arr[:, 1]]) out["endpoint_corner_mean"] = float(e_endpoint_min.mean()) out["endpoint_corner_min"] = float(e_endpoint_min.min()) return out def _stage1_displacement_features( verts: np.ndarray, s1_runs: Optional[List[Tuple[np.ndarray, Any, Any]]], ) -> Tuple[float, float, float]: """Distance from candidate vertices to nearest stage-1 prediction. Returns ``(disp_mean, disp_max, disp_far_frac)``, where ``disp_far_frac`` is the fraction of candidate vertices > 0.5 m from any stage-1 vertex. """ if not s1_runs or verts.size == 0: return 0.0, 0.0, 0.0 cv = np.asarray(verts, dtype=np.float64).reshape(-1, 3) if cv.shape[0] == 0: return 0.0, 0.0, 0.0 s1_v: List[np.ndarray] = [] for entry in s1_runs: rv = entry[0] if isinstance(entry, (tuple, list)) else entry rv = np.asarray(rv, dtype=np.float64).reshape(-1, 3) if rv.shape[0] > 0: s1_v.append(rv) if not s1_v: return 0.0, 0.0, 0.0 pool = np.concatenate(s1_v, axis=0) D = np.linalg.norm(cv[:, None, :] - pool[None, :, :], axis=-1) d_min = D.min(axis=1) return (float(d_min.mean()), float(d_min.max()), float(np.mean(d_min > 0.5))) def _connectivity_features( verts: np.ndarray, edges: List[Tuple[int, int]], ) -> Tuple[float, float, float]: """``(n_components_norm, cycle_rank_norm, degree_entropy)``. * ``n_components_norm`` = (n_components - 1) / max(1, Nv) — 0 when fully connected, grows with disconnected fragments. * ``cycle_rank_norm`` = (Ne - Nv + n_components) / max(1, Nv) — Betti-1 density; very low (tree-like) or very high (over-connected) are both suspicious. * ``degree_entropy`` = Shannon entropy of the degree distribution in nats. """ v = np.asarray(verts, dtype=np.float64).reshape(-1, 3) Nv = v.shape[0] if Nv == 0 or not edges: return 0.0, 0.0, 0.0 e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2) Ne = e_arr.shape[0] adj: List[List[int]] = [[] for _ in range(Nv)] for a, b in e_arr: if 0 <= a < Nv and 0 <= b < Nv and a != b: adj[a].append(b); adj[b].append(a) visited = [False] * Nv n_components = 0 for s in range(Nv): if visited[s]: continue n_components += 1 stack = [s] while stack: u = stack.pop() if visited[u]: continue visited[u] = True for w in adj[u]: if not visited[w]: stack.append(w) cycle_rank = Ne - Nv + n_components deg = np.array([len(a) for a in adj], dtype=np.float64) p = np.bincount(deg.astype(np.int64)) p = p / p.sum() if p.sum() > 0 else p p = p[p > 0] deg_entropy = float(-(p * np.log(p)).sum()) if p.size > 0 else 0.0 return ( float((n_components - 1) / max(1, Nv)), float(cycle_rank / max(1, Nv)), deg_entropy, ) def _symmetry_residual( verts: np.ndarray, edges: List[Tuple[int, int]], ) -> float: """Best reflection-plane residual normalised by mean edge length. Computes PCA on the candidate vertices, tries reflecting across each of the three principal planes (normal = each eigenvector through the centroid), Hungarian-matches reflected vs. original verts, and returns the minimum mean matched distance / mean edge length. Lower = more symmetric. 0 when the wireframe has fewer than 4 vertices or any edge. """ from scipy.optimize import linear_sum_assignment v = np.asarray(verts, dtype=np.float64).reshape(-1, 3) Nv = v.shape[0] if Nv < 4 or not edges: return 0.0 e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2) lengths = np.linalg.norm(v[e_arr[:, 0]] - v[e_arr[:, 1]], axis=1) mean_len = float(lengths.mean()) if lengths.size > 0 else 0.0 if mean_len < 1e-6: return 0.0 centroid = v.mean(axis=0) Q = v - centroid try: _, S, Vt = np.linalg.svd(Q, full_matrices=False) except np.linalg.LinAlgError: return 0.0 axes = Vt # rows are unit principal directions best = float("inf") for k in range(min(3, axes.shape[0])): n = axes[k] # Reflect Q across plane through centroid with normal n. ref = Q - 2.0 * (Q @ n)[:, None] * n[None, :] D = np.linalg.norm(Q[:, None, :] - ref[None, :, :], axis=-1) row, col = linear_sum_assignment(D) matched_d = D[row, col] cost = float(matched_d.mean()) if cost < best: best = cost if not np.isfinite(best): return 0.0 return float(best / mean_len) def _per_edge_class_consistency( verts: np.ndarray, edges: List[Tuple[int, int]], views: List[Dict[str, Any]], n_edge_samples: int = 8, ) -> Tuple[float, float]: """Gestalt-class purity per edge, aggregated to candidate-level stats. For each edge we sample interior points, project to each view, and look up the gestalt-class id at the pixel. Across all (sample × view) hits we take the modal class and define edge-purity as the fraction of hits belonging to that mode. Returns ``(mean_purity, high_purity_frac)`` where ``high_purity_frac = mean(purity > 0.5)``. Both are 0 with no views or edges. """ if not views or not edges or verts.size == 0: return 0.0, 0.0 v = np.asarray(verts, dtype=np.float64).reshape(-1, 3) e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2) Ne = e_arr.shape[0] ts = np.linspace(0.0, 1.0, n_edge_samples + 2)[1:-1] a = v[e_arr[:, 0]]; b = v[e_arr[:, 1]] samples = a[:, None, :] + (b - a)[:, None, :] * ts[None, :, None] sflat = samples.reshape(-1, 3) # (Ne*S, V) classes; -1 = no hit cls_per_view: List[np.ndarray] = [] for view in views: R = view["R"]; t = view["t"] p_cam = sflat @ R.T + t z = p_cam[:, 2] valid = z > 1e-6 zsafe = np.where(valid, z, 1.0) u_proj = p_cam[:, 0] / zsafe * view["fx"] + view["cx"] v_proj = p_cam[:, 1] / zsafe * view["fy"] + view["cy"] sx = view["W"] / view["cam_w"]; sy = view["H"] / view["cam_h"] ui = np.rint(u_proj * sx).astype(np.int64) vi = np.rint(v_proj * sy).astype(np.int64) in_view = (valid & (ui >= 0) & (ui < view["W"]) & (vi >= 0) & (vi < view["H"])) cls = -np.ones(sflat.shape[0], dtype=np.int64) if in_view.any(): idx = np.where(in_view)[0] cls[idx] = view["ids"][vi[idx], ui[idx]] # Only keep edge-class hits (other classes are background-ish noise) is_edge_cls = np.isin(cls, _EDGE_GESTALT_IDS) cls = np.where(is_edge_cls, cls, -1) cls_per_view.append(cls) if not cls_per_view: return 0.0, 0.0 cls_stack = np.stack(cls_per_view, axis=0).reshape( len(views), Ne, n_edge_samples) # For each edge: flatten (V*S) hits, drop -1, compute mode purity. purities = np.zeros(Ne, dtype=np.float64) for e_i in range(Ne): hits = cls_stack[:, e_i, :].reshape(-1) hits = hits[hits >= 0] if hits.size == 0: purities[e_i] = 0.0 continue counts = np.bincount(hits.astype(np.int64)) purities[e_i] = float(counts.max()) / float(hits.size) return float(purities.mean()), float(np.mean(purities > 0.5)) def _structural_features( verts: np.ndarray, edges: List[Tuple[int, int]], ) -> Tuple[float, float]: """Cheap structural priors: degree-1 vertex fraction + edge-length CV.""" v = np.asarray(verts, dtype=np.float64).reshape(-1, 3) Nv = v.shape[0] if Nv == 0 or not edges: return 0.0, 0.0 e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2) deg = np.zeros(Nv, dtype=np.int64) for a, b in e_arr: if 0 <= a < Nv: deg[a] += 1 if 0 <= b < Nv: deg[b] += 1 degree1_frac = float(np.mean(deg == 1)) lengths = np.linalg.norm(v[e_arr[:, 0]] - v[e_arr[:, 1]], axis=1) mean_len = float(lengths.mean()) if lengths.size > 0 else 0.0 cv = float(lengths.std() / mean_len) if mean_len > 1e-6 else 0.0 return degree1_frac, cv def _coplanarity_score( verts: np.ndarray, edges: List[Tuple[int, int]], max_cycles: int = 64, ) -> float: """Median plane-fit residual of 4-cycles, mapped to [0, 1] (higher = flatter). Buildings are mostly composed of planar quads (roof faces, wall panels). For every length-4 cycle in the wireframe we compute the third singular value of the centred (4,3) vertex matrix — the out-of-plane spread — normalised by the cycle's average edge length. The median residual is then mapped to ``exp(-5*r)`` so the feature is in [0, 1] with 1 = perfectly planar. Returns 0 when no cycles exist. """ Nv = verts.shape[0] if Nv < 4 or not edges: return 0.0 adj = [set() for _ in range(Nv)] for a, b in edges: a, b = int(a), int(b) if 0 <= a < Nv and 0 <= b < Nv and a != b: adj[a].add(b); adj[b].add(a) cycles: List[Tuple[int, int, int, int]] = [] for u in range(Nv): for v_idx in range(u + 1, Nv): common = adj[u] & adj[v_idx] if len(common) < 2: continue comm = sorted(c for c in common if c != u and c != v_idx) for i in range(len(comm)): for j in range(i + 1, len(comm)): cycles.append((u, comm[i], v_idx, comm[j])) if len(cycles) >= max_cycles: break if len(cycles) >= max_cycles: break if len(cycles) >= max_cycles: break if len(cycles) >= max_cycles: break if not cycles: return 0.0 vals = [] for cyc in cycles: pts = verts[list(cyc)] centroid = pts.mean(axis=0) Q = pts - centroid try: S = np.linalg.svd(Q, compute_uv=False) except np.linalg.LinAlgError: continue if S.size < 3: continue dev = float(S[-1]) perim = 0.0 for k in range(4): a, b = cyc[k], cyc[(k + 1) % 4] perim += float(np.linalg.norm(verts[a] - verts[b])) mean_len = perim / 4.0 if mean_len < 1e-6: continue vals.append(dev / mean_len) if not vals: return 0.0 return float(np.exp(-5.0 * float(np.median(vals)))) def _junction_angle_dev( verts: np.ndarray, edges: List[Tuple[int, int]], ) -> float: """Mean angular deviation of edge-pair junctions from {90°, 180°} in radians. At each vertex with degree >= 2, we compute the angle between every pair of incident edges and take the smaller deviation to 90° or 180° (the two angles that dominate roof/wall geometry). Lower is better; 0 means every junction sits exactly at a right or straight angle. """ Nv = verts.shape[0] if Nv == 0 or not edges: return 0.0 nbrs: List[List[int]] = [[] for _ in range(Nv)] for a, b in edges: a, b = int(a), int(b) if 0 <= a < Nv and 0 <= b < Nv and a != b: nbrs[a].append(b); nbrs[b].append(a) devs: List[float] = [] for u in range(Nv): if len(nbrs[u]) < 2: continue dirs: List[np.ndarray] = [] for v_idx in nbrs[u]: d = verts[v_idx] - verts[u] n = float(np.linalg.norm(d)) if n > 1e-6: dirs.append(d / n) for i in range(len(dirs)): for j in range(i + 1, len(dirs)): c = float(np.clip(np.dot(dirs[i], dirs[j]), -1.0, 1.0)) ang = float(np.arccos(c)) devs.append(min(abs(ang - np.pi / 2), abs(ang - np.pi))) if not devs: return 0.0 return float(np.mean(devs)) def _principal_axis_features( verts: np.ndarray, edges: List[Tuple[int, int]], cos_thresh: float = 0.9659, # cos(15°) ) -> Tuple[float, float]: """Manhattan structure features from per-edge direction PCA. Returns ``(manhattan_frac, top1_frac)``: * ``manhattan_frac`` — fraction of edges within 15° of any of the top-3 eigenvectors of the symmetric second-moment matrix ``sum_i d_i d_i^T`` (antipodal-symmetric: signs don't matter). * ``top1_frac`` — same, but only against the *dominant* eigenvector; a proxy for "fraction of edges aligned with the building's principal axis" (often vertical in COLMAP world). """ if not edges or verts.shape[0] == 0: return 0.0, 0.0 e_arr = np.asarray(edges, dtype=np.int64).reshape(-1, 2) if e_arr.shape[0] == 0: return 0.0, 0.0 d = verts[e_arr[:, 1]] - verts[e_arr[:, 0]] nrm = np.linalg.norm(d, axis=1) valid = nrm > 1e-6 if not valid.any(): return 0.0, 0.0 u = d[valid] / nrm[valid, None] M = u.T @ u try: eigvals, eigvecs = np.linalg.eigh(M) except np.linalg.LinAlgError: return 0.0, 0.0 order = np.argsort(-eigvals) eigvecs = eigvecs[:, order] dots = np.abs(u @ eigvecs) max_any = dots.max(axis=1) max_top1 = dots[:, 0] return (float(np.mean(max_any >= cos_thresh)), float(np.mean(max_top1 >= cos_thresh))) _RANKER_FEATURE_NAMES = [ "n_v", "n_e", "edge_density", "logit_mean", "logit_std", "logit_min", "logit_max", "vertex_support_mean", "vertex_support_std", "vertex_support_min", "edge_support_mean", "edge_support_max", "centrality_mean", "centrality_min", "centrality_rank", "is_raw_member", "is_medoid_source", "is_conf_source", "is_refined", "uses_edge_vote", "uses_add_edges", "uses_topk", "source_idx", # --- 2D reprojection + structural priors (added) --- "vert_reproj_mean", "vert_reproj_min", "edge_reproj_mean", "edge_reproj_min", "degree1_frac", "edge_length_cv", # --- geometric priors (added) --- "coplanarity_score", "junction_angle_dev", "manhattan_frac", "principal_axis_align", # --- v5 features (selected from v4 ablation: only those that survived # 5-fold CV with positive lift in addition to v3) --- "s1_disp_far_frac", "degree_entropy", "symmetry_residual", ] def _edge_key_set(edges: List[Tuple[int, int]]) -> set: out = set() for a, b in edges: ia, ib = int(a), int(b) if ia == ib: continue out.add((min(ia, ib), max(ia, ib))) return out def _candidate_logits_from_source( runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], source_idx: int, keep_idx: Optional[np.ndarray], n_verts: int, ) -> np.ndarray: if source_idx < 0 or source_idx >= len(runs): return np.zeros((n_verts,), dtype=np.float64) lg = np.asarray(runs[source_idx][2], dtype=np.float64).reshape(-1) if keep_idx is not None and lg.size > 0: ki = np.asarray(keep_idx, dtype=np.int64).reshape(-1) valid = (ki >= 0) & (ki < lg.size) if valid.all() and ki.size == n_verts: return lg[ki] if lg.size == n_verts: return lg if lg.size == 0: return np.zeros((n_verts,), dtype=np.float64) return np.full((n_verts,), float(np.nanmean(lg)), dtype=np.float64) def _candidate_support_features( verts: np.ndarray, edges: List[Tuple[int, int]], runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], match_tau_m: float, ) -> Tuple[float, float, float, float, float]: """Agreement features for one candidate against the ensemble members.""" from scipy.optimize import linear_sum_assignment v = np.asarray(verts, dtype=np.float64).reshape(-1, 3) Nv = v.shape[0] if Nv == 0 or not runs: return 0.0, 0.0, 0.0, 0.0, 0.0 v_support = np.zeros((Nv,), dtype=np.float64) e_keys = _edge_key_set(edges) e_votes: Dict[Tuple[int, int], int] = {k: 0 for k in e_keys} for rv_raw, re_raw, _rlg in runs: rv = np.asarray(rv_raw, dtype=np.float64).reshape(-1, 3) if rv.shape[0] == 0: continue D = np.linalg.norm(v[:, None, :] - rv[None, :, :], axis=-1) D_capped = np.minimum(D, match_tau_m) row_ind, col_ind = linear_sum_assignment(D_capped) r_to_c = -np.ones(rv.shape[0], dtype=np.int64) for r, c in zip(row_ind, col_ind): if D[r, c] <= match_tau_m: v_support[r] += 1.0 r_to_c[c] = r if e_keys: seen = set() for a, b in re_raw: ia, ib = int(a), int(b) if (ia == ib or ia < 0 or ib < 0 or ia >= r_to_c.size or ib >= r_to_c.size): continue ca, cb = int(r_to_c[ia]), int(r_to_c[ib]) if ca < 0 or cb < 0 or ca == cb: continue key = (min(ca, cb), max(ca, cb)) if key in e_votes: seen.add(key) for key in seen: e_votes[key] += 1 denom = float(max(1, len(runs))) v_frac = v_support / denom if e_votes: e_frac = np.asarray(list(e_votes.values()), dtype=np.float64) / denom e_mean = float(e_frac.mean()) e_max = float(e_frac.max()) else: e_mean = 0.0 e_max = 0.0 return ( float(v_frac.mean()), float(v_frac.std()), float(v_frac.min()), e_mean, e_max, ) def _candidate_feature_vector( name: str, verts: np.ndarray, edges: List[Tuple[int, int]], logits: np.ndarray, runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], match_tau_m: float, source_idx: int, medoid_idx: int, confidence_idx: int, ranking: List[int], is_raw_member: bool, is_refined: bool, uses_edge_vote: bool, uses_add_edges: bool, uses_topk: bool, reproj_views: Optional[List[Dict[str, Any]]] = None, s1_runs: Optional[List[Tuple[np.ndarray, Any, Any]]] = None, ) -> np.ndarray: v = np.asarray(verts, dtype=np.float64).reshape(-1, 3) e = list(edges) lg = np.asarray(logits, dtype=np.float64).reshape(-1) lg = lg[np.isfinite(lg)] if lg.size == 0: lg = np.zeros((1,), dtype=np.float64) n_v = float(v.shape[0]) n_e = float(len(e)) max_edges = max(1.0, n_v * max(0.0, n_v - 1.0) / 2.0) v_sup_mean, v_sup_std, v_sup_min, e_sup_mean, e_sup_max = ( _candidate_support_features(v, e, runs, match_tau_m) ) dists = [] for rv, re, _rl in runs: try: dists.append(_wireframe_pair_distance( v, e, np.asarray(rv, dtype=np.float64), list(re), match_tau_m=match_tau_m, )) except Exception: continue d_arr = np.asarray([d for d in dists if np.isfinite(d)], dtype=np.float64) if d_arr.size == 0: centrality_mean = 10.0 centrality_min = 10.0 else: centrality_mean = float(d_arr.mean()) centrality_min = float(d_arr.min()) rank_pos = float(len(ranking)) if source_idx in ranking: rank_pos = float(ranking.index(source_idx)) if reproj_views: vr_mean, vr_min, er_mean, er_min = _reprojection_features( v, e, reproj_views, ) else: vr_mean = vr_min = er_mean = er_min = 0.0 deg1_frac, edge_len_cv = _structural_features(v, e) coplanarity = _coplanarity_score(v, e) junction_dev = _junction_angle_dev(v, e) manhattan_frac, princ_axis_align = _principal_axis_features(v, e) _, _, s1_disp_far = _stage1_displacement_features(v, s1_runs) _, _, deg_ent = _connectivity_features(v, e) sym_res = _symmetry_residual(v, e) feats = np.asarray([ n_v / 64.0, n_e / 128.0, n_e / max_edges, float(lg.mean()), float(lg.std()), float(lg.min()), float(lg.max()), v_sup_mean, v_sup_std, v_sup_min, e_sup_mean, e_sup_max, centrality_mean, centrality_min, rank_pos / float(max(1, len(ranking))), 1.0 if is_raw_member else 0.0, 1.0 if source_idx == medoid_idx else 0.0, 1.0 if source_idx == confidence_idx else 0.0, 1.0 if is_refined else 0.0, 1.0 if uses_edge_vote else 0.0, 1.0 if uses_add_edges else 0.0, 1.0 if uses_topk else 0.0, float(source_idx) / float(max(1, len(runs) - 1)), vr_mean, vr_min, er_mean, er_min, deg1_frac, edge_len_cv, coplanarity, junction_dev, manhattan_frac, princ_axis_align, # --- v5 features --- s1_disp_far, deg_ent, sym_res, ], dtype=np.float64) if feats.shape[0] != len(_RANKER_FEATURE_NAMES): raise RuntimeError(f"ranker feature mismatch for {name}") return feats def _ensemble_heuristic_score(features: np.ndarray) -> float: f = {k: float(v) for k, v in zip(_RANKER_FEATURE_NAMES, features)} return ( 1.20 * f["vertex_support_mean"] + 0.55 * f["edge_support_mean"] + 0.08 * f["logit_mean"] - 0.30 * f["centrality_mean"] - 0.08 * abs(f["n_v"] - 0.33) - 0.04 * abs(f["n_e"] - 0.18) + 0.04 * f["is_refined"] + 0.02 * f["is_medoid_source"] - 0.03 * f["uses_edge_vote"] ) def _score_ranker_candidate( features: np.ndarray, ranker: Optional[Dict[str, np.ndarray]], ) -> float: if ranker is None: return _ensemble_heuristic_score(features) mean = np.asarray(ranker["mean"], dtype=np.float64) scale = np.asarray(ranker["scale"], dtype=np.float64) weights = np.asarray(ranker["weights"], dtype=np.float64) x = (features - mean) / np.maximum(scale, 1e-8) return float(np.dot(x, weights[1:]) + weights[0]) def _fuse_topk_candidates( candidates: List[Dict[str, Any]], scores: np.ndarray, k: int, match_tau_m: float, edge_vote_frac: float, ) -> Tuple[np.ndarray, List[Tuple[int, int]], int]: """Fuse the top-K ranker-scored candidates into one wireframe. Strategy mirrors ``medoid_mean_add`` but applies it to the top-K *ranker-quality-filtered* candidates rather than all ensemble members: 1. Sort candidates by ``scores`` descending; keep ``top_k = min(k, len)``. 2. Pivot = top-1 candidate. For each pivot vertex, Hungarian-match to each of the other top-K candidates' vertices (cap τ = ``match_tau_m``) and average the positions of matched correspondents. 3. Union pivot's edges with edges voted by at least ``edge_vote_frac`` of the top-K (counted on matched vertex pairs). Returns ``(verts, edges, best_idx)`` where ``best_idx`` is the position of the top-1 candidate in the input list (for diagnostics). """ scores_arr = np.asarray(scores, dtype=np.float64).reshape(-1) n = scores_arr.size if n == 0: return np.zeros((0, 3), dtype=np.float64), [], 0 k_eff = max(1, min(int(k), n)) order = np.argsort(-scores_arr)[:k_eff].astype(np.int64) best_idx = int(order[0]) if k_eff == 1: c = candidates[best_idx] return (np.asarray(c["verts"], dtype=np.float64), list(c["edges"]), best_idx) runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]] = [] for j in order: c = candidates[int(j)] v = np.asarray(c["verts"], dtype=np.float64).reshape(-1, 3) e = list(c["edges"]) lg_src = c.get("logits", None) lg = (np.asarray(lg_src, dtype=np.float64).reshape(-1) if lg_src is not None and len(np.asarray(lg_src).ravel()) == v.shape[0] else np.zeros((v.shape[0],), dtype=np.float64)) runs.append((v, e, lg)) refined_v, refined_e, keep_idx = _refine_positions( runs, picked_idx=0, match_tau_m=match_tau_m, aggregator="mean", member_indices=None, vertex_vote_frac=0.0, ) voted = _vote_edges( runs, picked_idx=0, match_tau_m=match_tau_m, vote_frac=edge_vote_frac, member_indices=None, keep_idx=keep_idx, ) edge_set = set() for (a, b) in refined_e: ia, ib = int(a), int(b) if ia == ib: continue edge_set.add((min(ia, ib), max(ia, ib))) for (a, b) in voted: ia, ib = int(a), int(b) if ia == ib: continue edge_set.add((min(ia, ib), max(ia, ib))) return refined_v, sorted(edge_set), best_idx def _load_ensemble_ranker(path: Optional[str]) -> Optional[Dict[str, np.ndarray]]: if not path: return None p = Path(path) if not p.exists(): print(f"[ranker] {p} not found; using built-in heuristic scorer", flush=True) return None data = np.load(p, allow_pickle=False) names = [str(x) for x in data["feature_names"].tolist()] if names != _RANKER_FEATURE_NAMES: print( f"[ranker] feature-name mismatch in {p} " f"(have {len(names)}-d, expected {len(_RANKER_FEATURE_NAMES)}-d); " f"falling back to built-in heuristic scorer", flush=True) return None ranker = { "mean": np.asarray(data["mean"], dtype=np.float64), "scale": np.asarray(data["scale"], dtype=np.float64), "weights": np.asarray(data["weights"], dtype=np.float64), } print(f"[ranker] loaded {p}", flush=True) return ranker def _build_ranker_candidates( runs: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], match_tau_m: float, edge_vote_frac: float, ranking: Optional[List[int]] = None, medoid_idx: Optional[int] = None, confidence_idx: Optional[int] = None, reproj_views: Optional[List[Dict[str, Any]]] = None, s1_runs: Optional[List[Tuple[np.ndarray, Any, Any]]] = None, ) -> List[Dict[str, Any]]: """Produce candidate wireframes for learned/heuristic ensemble selection.""" if not runs: return [] if ranking is None: ranking = _medoid_rank(runs, match_tau_m=match_tau_m) medoid_idx = int(medoid_idx if medoid_idx is not None else (ranking[0] if ranking else 0)) confidence_idx = int(confidence_idx if confidence_idx is not None else _confidence_select_idx(runs)) candidates: List[Dict[str, Any]] = [] def add_candidate( name: str, verts: np.ndarray, edges: List[Tuple[int, int]], source_idx: int, keep_idx: Optional[np.ndarray], is_raw_member: bool, is_refined: bool, uses_edge_vote: bool, uses_add_edges: bool, uses_topk: bool, ) -> None: v = np.asarray(verts, dtype=np.float64).reshape(-1, 3) e = [(int(a), int(b)) for a, b in edges if int(a) != int(b)] if v.shape[0] == 0 or len(e) == 0: return logits = _candidate_logits_from_source(runs, source_idx, keep_idx, v.shape[0]) features = _candidate_feature_vector( name=name, verts=v, edges=e, logits=logits, runs=runs, match_tau_m=match_tau_m, source_idx=source_idx, medoid_idx=medoid_idx, confidence_idx=confidence_idx, ranking=ranking, is_raw_member=is_raw_member, is_refined=is_refined, uses_edge_vote=uses_edge_vote, uses_add_edges=uses_add_edges, uses_topk=uses_topk, reproj_views=reproj_views, s1_runs=s1_runs, ) candidates.append({ "name": name, "verts": v, "edges": e, "logits": logits, "features": features, "source_idx": int(source_idx), }) for i, (v, e, _lg) in enumerate(runs): add_candidate( name=f"member_{i:02d}", verts=v, edges=list(e), source_idx=i, keep_idx=None, is_raw_member=True, is_refined=False, uses_edge_vote=False, uses_add_edges=False, uses_topk=False, ) def refined_candidates(source_idx: int, label: str, member_indices: Optional[List[int]]) -> None: uses_topk = member_indices is not None rv, re, keep_idx = _refine_positions( runs, source_idx, match_tau_m=match_tau_m, aggregator="mean", member_indices=member_indices, vertex_vote_frac=0.0, ) add_candidate( name=f"{label}_mean_picked" + ("_topk" if uses_topk else ""), verts=rv, edges=re, source_idx=source_idx, keep_idx=keep_idx, is_raw_member=False, is_refined=True, uses_edge_vote=False, uses_add_edges=False, uses_topk=uses_topk, ) voted = _vote_edges( runs, source_idx, match_tau_m=match_tau_m, vote_frac=edge_vote_frac, member_indices=member_indices, keep_idx=keep_idx, ) add_candidate( name=f"{label}_mean_voted" + ("_topk" if uses_topk else ""), verts=rv, edges=voted, source_idx=source_idx, keep_idx=keep_idx, is_raw_member=False, is_refined=True, uses_edge_vote=True, uses_add_edges=False, uses_topk=uses_topk, ) add_edges = sorted(_edge_key_set(re) | _edge_key_set(voted)) add_candidate( name=f"{label}_mean_add" + ("_topk" if uses_topk else ""), verts=rv, edges=add_edges, source_idx=source_idx, keep_idx=keep_idx, is_raw_member=False, is_refined=True, uses_edge_vote=True, uses_add_edges=True, uses_topk=uses_topk, ) refined_candidates(medoid_idx, "medoid", None) if confidence_idx != medoid_idx: refined_candidates(confidence_idx, "confidence", None) for k in (5, 8): if ranking and len(ranking) >= k: members = list(ranking[:k]) if medoid_idx not in members: members.append(medoid_idx) refined_candidates(medoid_idx, f"medoid_top{k}", members) return candidates def train_ensemble_ranker_jsonl(jsonl_path: str, out_path: str, l2: float = 1e-2) -> None: """Fit a tiny ridge ranker from within-sample candidate comparisons. Absolute HSS is dominated by sample difficulty ("easy building" vs. "hard building"), but inference only needs to choose among candidates for the same sample. Training on pairwise feature differences cancels that nuisance and directly optimizes the ordering we care about. """ rows_by_sample: Dict[str, List[Tuple[np.ndarray, float]]] = {} with open(jsonl_path, "r") as f: for line in f: line = line.strip() if not line: continue row = json.loads(line) feats = np.asarray(row.get("features", []), dtype=np.float64) if feats.shape[0] != len(_RANKER_FEATURE_NAMES): continue y = float(row.get("hss", 0.0)) if not np.isfinite(y) or not np.all(np.isfinite(feats)): continue sid = str(row.get("order_id", row.get("sample_index", ""))) rows_by_sample.setdefault(sid, []).append((feats, y)) xs = [feat for rows in rows_by_sample.values() for feat, _y in rows] if len(xs) < 10 or len(rows_by_sample) < 2: raise SystemExit( f"[ranker/train] need >=10 rows and >=2 samples, got " f"{len(xs)} rows / {len(rows_by_sample)} samples") X = np.stack(xs, axis=0) mean = X.mean(axis=0) scale = X.std(axis=0) scale[scale < 1e-8] = 1.0 pair_x: List[np.ndarray] = [] pair_y: List[float] = [] for rows in rows_by_sample.values(): if len(rows) < 2: continue feats = [(feat - mean) / scale for feat, _y in rows] ys = [float(y) for _feat, y in rows] for i in range(len(rows)): for j in range(i + 1, len(rows)): dy = ys[i] - ys[j] if abs(dy) < 1e-9: continue pair_x.append(feats[i] - feats[j]) pair_y.append(dy) if len(pair_x) < 10: raise SystemExit(f"[ranker/train] need >=10 informative pairs, got {len(pair_x)}") DX = np.stack(pair_x, axis=0) dy = np.asarray(pair_y, dtype=np.float64) reg = np.eye(DX.shape[1], dtype=np.float64) * float(l2) w = np.linalg.solve(DX.T @ DX + reg, DX.T @ dy) weights = np.concatenate([[0.0], w], axis=0) pred = DX @ w mse = float(np.mean((pred - dy) ** 2)) corr = float(np.corrcoef(pred, dy)[0, 1]) if pred.std() > 0 and dy.std() > 0 else 0.0 # Grouped oracle estimate: how good would this ranker be on the same dump? # This is optimistic when train==eval, but useful as a smoke test. picked = [] oracle = [] for rows in rows_by_sample.values(): scored = [] for feat, y in rows: x = np.concatenate([[1.0], (feat - mean) / scale]) scored.append((float(np.dot(x, weights)), float(y))) picked.append(max(scored, key=lambda t: t[0])[1]) oracle.append(max(scored, key=lambda t: t[1])[1]) np.savez( out_path, feature_names=np.asarray(_RANKER_FEATURE_NAMES, dtype=" Tuple[Optional[np.ndarray], Optional[List[Tuple[int, int]]], Optional[np.ndarray], bool]: """Build stage-2 scene from a single stage-1 prediction and run stage-2 refinement. Returns ``(verts_world, edges, valid_logits, hull_valid)``. ``valid_logits`` is a (M,) float32 array of raw validity logits for the kept vertices (used by confidence-mode ensemble selection). When the hull is empty or stage-2 produces nothing, returns ``(None, None, None, False)`` so the caller can fall back to stage-1.""" def _sync(): if device.type == "cuda": torch.cuda.synchronize() rng = np.random.default_rng(_seed_from_order_id(order_id)) t0 = time.perf_counter() s2_row = build_stage2_scene( raw_sample=raw_sample, colmap_rec=colmap_rec, pred_verts=pred_verts, pred_edges=pred_edges, above_m=stage2_cfg.above_m, below_m=stage2_cfg.below_m, side_m=stage2_cfg.side_m, n_col_oversample=stage2_cfg.n_col_oversample, n_dep_oversample=stage2_cfg.n_dep_oversample, n_depth_per_image_cap=stage2_cfg.n_depth_per_image_cap, sampling_mode=("voxel" if stage2_cfg.sampling == "voxel" else "random"), density_voxel_size_m=stage2_cfg.density_voxel_size_m, density_kernel_radius=stage2_cfg.density_kernel_radius, density_kernel_axis=stage2_cfg.density_kernel_axis, density_response_power=stage2_cfg.density_response_power, density_planarity_suppression=stage2_cfg.density_planarity_suppression, density_planarity_radius=stage2_cfg.density_planarity_radius, density_planarity_min_points=stage2_cfg.density_planarity_min_points, density_min_per_voxel=stage2_cfg.density_min_per_voxel, rng=rng, cache=preprocess_cache, ) if timings_accum is not None: timings_accum["s2_build"] = timings_accum.get("s2_build", 0.0) + (time.perf_counter() - t0) t0 = time.perf_counter() s2_sample = stage2_row_to_sample( s2_row, n_pts=stage2_cfg.n_pts, k_verts=stage2_cfg.k_verts, augment=False, flip=False, yaw=False, jitter_sigma=0.0, pre_subsample=(stage2_cfg.sampling != "fps"), ) if timings_accum is not None: timings_accum["s2_transform"] = timings_accum.get("s2_transform", 0.0) + (time.perf_counter() - t0) batch2 = _stage2_sample_to_batch(s2_sample, device) if stage2_cfg.sampling == "fps": seed = torch.zeros(1, device=device, dtype=torch.long) batch2 = fps_subsample_stage2_batch( batch2, n_col=stage2_cfg.n_pts // 2, n_dep=stage2_cfg.n_pts - stage2_cfg.n_pts // 2, seed_per_sample=seed, renormalize_bbox=True, max_exact_iters_per_provenance=stage2_cfg.fps_max_exact_iters, ) s2_model = stage2_cfg.model n_steps2 = int(stage2_cfg.n_sample_steps) vth2 = float(stage2_cfg.validity_thresh) K2 = s2_model.denoiser.k_verts B2 = 1 step_dt2 = 1.0 / max(1, n_steps2) with torch.inference_mode(): t0 = time.perf_counter() scene_feats2, scene_xyz_norm2 = s2_model._encode_scene(batch2) _sync() if timings_accum is not None: timings_accum["s2_enc"] = timings_accum.get("s2_enc", 0.0) + (time.perf_counter() - t0) x = s2_model._init_x0(batch2, K2, device, B2) last_logit = torch.zeros(B2, K2, device=device) last_edge_logit = torch.zeros(B2, K2, K2, device=device) last_k = max(1, int(getattr(stage2_cfg, "ensemble_stage2_last_k", 1))) start_avg = n_steps2 - last_k x_sum: Optional[torch.Tensor] = None logit_sum: Optional[torch.Tensor] = None edge_sum: Optional[torch.Tensor] = None n_accum = 0 t0 = time.perf_counter() for s_idx in range(n_steps2): t = torch.full((B2,), s_idx * step_dt2, device=device) v, last_logit, last_edge_logit = s2_model.denoiser( x, t, scene_feats2, scene_xyz_norm2, ) x = x + step_dt2 * v if last_k > 1 and s_idx >= start_avg: if x_sum is None: x_sum = x.clone() logit_sum = last_logit.clone() edge_sum = last_edge_logit.clone() else: x_sum.add_(x) logit_sum.add_(last_logit) edge_sum.add_(last_edge_logit) n_accum += 1 if x_sum is not None and n_accum > 0: inv = 1.0 / float(n_accum) x = x_sum * inv last_logit = logit_sum * inv last_edge_logit = edge_sum * inv _sync() if timings_accum is not None: timings_accum["s2_den"] = timings_accum.get("s2_den", 0.0) + (time.perf_counter() - t0) xyz_norm = x[0] valid = last_logit[0] > vth2 edge_logit_ = last_edge_logit[0].float() valid_idx = torch.nonzero(valid, as_tuple=False).flatten() xyz_valid = xyz_norm.index_select(0, valid_idx) center = batch2["bbox_center"][0].to(xyz_valid) scale = batch2["bbox_scale"][0].to(xyz_valid) bbox_R = batch2.get("bbox_R") if isinstance(bbox_R, torch.Tensor): R = bbox_R[0].to(xyz_valid) verts_world = (xyz_valid * scale) @ R + center else: verts_world = xyz_valid * scale + center s2_verts = verts_world.cpu().numpy() s2_logits = (last_logit[0].index_select(0, valid_idx) .detach().cpu().numpy().astype(np.float32)) s2_edges: List[Tuple[int, int]] = [] n_valid = int(valid_idx.numel()) if n_valid >= 2: sub = edge_logit_.index_select(0, valid_idx).index_select(1, valid_idx) tri = torch.triu(torch.ones(n_valid, n_valid, device=sub.device, dtype=torch.bool), diagonal=1) pairs = torch.nonzero((sub > 0.0) & tri, as_tuple=False) s2_edges = [(int(a), int(b)) for a, b in pairs.detach().cpu().tolist()] hull_valid = bool(s2_row.get("hull_valid", False)) if not hull_valid or len(s2_verts) == 0 or len(s2_edges) == 0: return None, None, None, hull_valid return np.asarray(s2_verts, dtype=float), s2_edges, s2_logits, hull_valid def _predict_union_hull_ensemble( raw_sample: Dict[str, Any], colmap_rec, s1_runs: List[Tuple[np.ndarray, np.ndarray, np.ndarray]], stage2_cfg: Stage2Config, device: torch.device, preprocess_cache: Optional[Dict[str, Any]], timings_accum: Optional[Dict[str, float]], order_id: str, ) -> Tuple[List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]], bool]: """Union-hull batched ensemble. Strategy: 1. Concatenate the N stage-1 vertex predictions into one cloud and build **one** stage-2 scene (single hull, single point cloud, single bbox). 2. Transform once → one shared sample. Encode once (B=1). 3. Build B=N denoiser batch sharing the scene but with per-slot ``init_verts`` derived from each stage-1 prediction (each slot starts its trajectory from a different stage-1 hypothesis). 4. Run the stage-2 Euler loop at B=N — broadcasts the shared encoded features via ``.expand`` (no extra GPU memory) for ``N`` parallel trajectories with diverse inits. This shares the heaviest stage-2 cost (encoder + build) across all N seeds, so total per-sample cost is ~1.0-1.2x single-run wall vs 1.5-1.8x for the per-seed strategy at N=5. Returns ``(seed_results, hull_valid)`` where ``seed_results`` is ``[(verts_world, edges, valid_logits)]`` per slot, in the same shape as ``_stage1_seed_ensemble``. Empty list on catastrophic failure; caller handles the fallback. """ if not s1_runs: return [], False # 1) Union of stage-1 vertex predictions for the hull. union_verts_parts = [] union_edges_parts = [] offset = 0 for v, e, _l in s1_runs: if len(v) == 0: continue union_verts_parts.append(np.asarray(v, dtype=np.float32).reshape(-1, 3)) if len(e) > 0: ea = np.asarray(e, dtype=np.int64).reshape(-1, 2) union_edges_parts.append(ea + offset) offset += len(v) if not union_verts_parts: return [], False union_verts = np.concatenate(union_verts_parts, axis=0) union_edges = (np.concatenate(union_edges_parts, axis=0) if union_edges_parts else np.zeros((0, 2), dtype=np.int64)) # 2) Build the single shared stage-2 scene. rng = np.random.default_rng(_seed_from_order_id(order_id)) t0 = time.perf_counter() s2_row = build_stage2_scene( raw_sample=raw_sample, colmap_rec=colmap_rec, pred_verts=union_verts, pred_edges=union_edges, above_m=stage2_cfg.above_m, below_m=stage2_cfg.below_m, side_m=stage2_cfg.side_m, n_col_oversample=stage2_cfg.n_col_oversample, n_dep_oversample=stage2_cfg.n_dep_oversample, n_depth_per_image_cap=stage2_cfg.n_depth_per_image_cap, sampling_mode=("voxel" if stage2_cfg.sampling == "voxel" else "random"), density_voxel_size_m=stage2_cfg.density_voxel_size_m, density_kernel_radius=stage2_cfg.density_kernel_radius, density_kernel_axis=stage2_cfg.density_kernel_axis, density_response_power=stage2_cfg.density_response_power, density_planarity_suppression=stage2_cfg.density_planarity_suppression, density_planarity_radius=stage2_cfg.density_planarity_radius, density_planarity_min_points=stage2_cfg.density_planarity_min_points, density_min_per_voxel=stage2_cfg.density_min_per_voxel, rng=rng, cache=preprocess_cache, ) if timings_accum is not None: timings_accum["s2_build"] = timings_accum.get("s2_build", 0.0) + (time.perf_counter() - t0) hull_valid = bool(s2_row.get("hull_valid", False)) if not hull_valid: return [], False # 3) Transform once → one shared sample (random-sampling path: scene gets # bbox-normalised here; init_verts is the field we'll override per slot). t0 = time.perf_counter() s2_sample = stage2_row_to_sample( s2_row, n_pts=stage2_cfg.n_pts, k_verts=stage2_cfg.k_verts, augment=False, flip=False, yaw=False, jitter_sigma=0.0, pre_subsample=(stage2_cfg.sampling != "fps"), ) if timings_accum is not None: timings_accum["s2_transform"] = timings_accum.get("s2_transform", 0.0) + (time.perf_counter() - t0) batch2_single = _stage2_sample_to_batch(s2_sample, device) if stage2_cfg.sampling == "fps": seed = torch.zeros(1, device=device, dtype=torch.long) batch2_single = fps_subsample_stage2_batch( batch2_single, n_col=stage2_cfg.n_pts // 2, n_dep=stage2_cfg.n_pts - stage2_cfg.n_pts // 2, seed_per_sample=seed, renormalize_bbox=True, max_exact_iters_per_provenance=stage2_cfg.fps_max_exact_iters, ) # 4) Per-slot init_verts: each slot starts from its own stage-1 prediction, # normalised into the shared bbox frame (same bbox the encoder sees). K = stage2_cfg.k_verts N = len(s1_runs) bbox_center_np = batch2_single["bbox_center"][0].detach().cpu().numpy().astype(np.float32) bbox_scale_val = float(batch2_single["bbox_scale"][0].detach().cpu().item()) init_padded_n = np.zeros((N, K, 3), dtype=np.float32) init_valid_n = np.zeros((N, K), dtype=bool) seed_root = _seed_from_order_id(order_id) for i, (verts_w, _e, _l) in enumerate(s1_runs): if len(verts_w) == 0: continue init_norm = ((np.asarray(verts_w, dtype=np.float32).reshape(-1, 3) - bbox_center_np) / max(bbox_scale_val, 1e-6)).astype(np.float32) rng_i = np.random.default_rng(seed_root + i + 1) ip, iv = _pad_init_verts(init_norm, K, rng=rng_i) init_padded_n[i] = ip init_valid_n[i] = iv # 5) Build B=N denoiser batch sharing every per-point field, overriding # only init_verts / init_verts_valid per slot. batch2_N = _replicate_batch(batch2_single, N) batch2_N["init_verts"] = torch.from_numpy(init_padded_n).to(device) batch2_N["init_verts_valid"] = torch.from_numpy(init_valid_n).to(device) # 6) Encoder once (B=1), broadcast features to B=N for the denoiser. s2_model = stage2_cfg.model n_steps2 = int(stage2_cfg.n_sample_steps) vth2 = float(stage2_cfg.validity_thresh) step_dt2 = 1.0 / max(1, n_steps2) def _sync(): if device.type == "cuda": torch.cuda.synchronize() with torch.inference_mode(): t0 = time.perf_counter() scene_feats_1, scene_xyz_1 = s2_model._encode_scene(batch2_single) scene_feats_N = scene_feats_1.expand(N, -1, -1).contiguous() scene_xyz_N = scene_xyz_1.expand(N, -1, -1).contiguous() _sync() if timings_accum is not None: timings_accum["s2_enc"] = timings_accum.get("s2_enc", 0.0) + (time.perf_counter() - t0) x = s2_model._init_x0(batch2_N, K, device, N) last_logit = torch.zeros(N, K, device=device) last_edge_logit = torch.zeros(N, K, K, device=device) last_k = max(1, int(getattr(stage2_cfg, "ensemble_stage2_last_k", 1))) start_avg = n_steps2 - last_k x_sum: Optional[torch.Tensor] = None logit_sum: Optional[torch.Tensor] = None edge_sum: Optional[torch.Tensor] = None n_accum = 0 t0 = time.perf_counter() for s_idx in range(n_steps2): t = torch.full((N,), s_idx * step_dt2, device=device) v, last_logit, last_edge_logit = s2_model.denoiser( x, t, scene_feats_N, scene_xyz_N, ) x = x + step_dt2 * v if last_k > 1 and s_idx >= start_avg: if x_sum is None: x_sum = x.clone() logit_sum = last_logit.clone() edge_sum = last_edge_logit.clone() else: x_sum.add_(x) logit_sum.add_(last_logit) edge_sum.add_(last_edge_logit) n_accum += 1 if x_sum is not None and n_accum > 0: inv = 1.0 / float(n_accum) x = x_sum * inv last_logit = logit_sum * inv last_edge_logit = edge_sum * inv _sync() if timings_accum is not None: timings_accum["s2_den"] = timings_accum.get("s2_den", 0.0) + (time.perf_counter() - t0) # 7) Extract per-slot predictions in world coords. results: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]] = [] bbox_R = batch2_N.get("bbox_R") for i in range(N): valid = last_logit[i] > vth2 idx = torch.nonzero(valid, as_tuple=False).flatten() x_valid = x[i].index_select(0, idx) center = batch2_N["bbox_center"][i].to(device=x_valid.device, dtype=torch.float32) scale = batch2_N["bbox_scale"][i].to(device=x_valid.device, dtype=torch.float32) if isinstance(bbox_R, torch.Tensor): R = bbox_R[i].to(device=x_valid.device, dtype=torch.float32) verts_world = (x_valid * scale) @ R + center else: verts_world = x_valid * scale + center verts_np = verts_world.detach().cpu().numpy().astype(np.float32).reshape(-1, 3) logits_kept = (last_logit[i].index_select(0, idx) .detach().cpu().numpy().astype(np.float32)) n_valid = int(idx.numel()) edges_i: List[Tuple[int, int]] = [] if n_valid >= 2: sub = last_edge_logit[i].float().index_select(0, idx).index_select(1, idx) tri = torch.triu(torch.ones(n_valid, n_valid, device=sub.device, dtype=torch.bool), diagonal=1) pairs = torch.nonzero((sub > 0.0) & tri, as_tuple=False) edges_i = [(int(a), int(b)) for a, b in pairs.detach().cpu().tolist()] results.append((verts_np, edges_i, logits_kept)) return results, True def predict_two_stage( raw_sample: Dict[str, Any], colmap_rec, stage1_scene: Dict[str, np.ndarray], stage1_model: WireframeDiffusion, stage1_n_steps: int, stage1_validity_thresh: float, stage2_cfg: Stage2Config, device: torch.device, preprocess_cache: Optional[Dict[str, Any]] = None, stats: Optional[Dict[str, int]] = None, timings: Optional[Dict[str, float]] = None, diagnostics: Optional[Dict[str, Any]] = None, ) -> Tuple[np.ndarray, List[Tuple[int, int]]]: """Run stage 1 (coarse) → build stage-2 scene from its prediction → run stage 2 (refinement). Returns the stage-2 verts + edges in world coords. When ``stage2_cfg.ensemble_n > 1``, runs N independent (stage-1, stage-2) trajectories with different random init x0 (stage-1 denoiser batched over seeds, stage-2 sequential since each seed produces its own hull) and fuses the N wireframes by vertex-merge consensus. If ``timings`` is provided, populates per-phase wallclock seconds summed across seeds: ``s2_build``, ``s2_transform``, ``s2_enc``, ``s2_den``. Stage-1 inference is NOT timed here because the sanity check inlines stage 1 with its own per-phase timing. """ order_id = str(raw_sample.get("order_id", "")) n_ensemble = max(1, int(getattr(stage2_cfg, "ensemble_n", 1))) strategy = str(getattr(stage2_cfg, "ensemble_strategy", "union_hull")) # 1) Stage-1 inference (GPU) — N seeds batched on one shared scene. batch1 = scene_to_batch(stage1_scene, device) with torch.inference_mode(): s1_runs = _stage1_seed_ensemble( stage1_model, batch1, n_steps=stage1_n_steps, validity_thresh=stage1_validity_thresh, n_seeds=n_ensemble, ) # 2) Stage-2 ensembling. # # - "union_hull": build ONE stage-2 scene from the union of all N # stage-1 vertex predictions, encode ONCE, then run a B=N denoiser # batch with per-slot init_verts from each stage-1 hypothesis. # Cheap (~1.0-1.2x single-run wall) and gives N stage-2 predictions # that share a common reference frame. # # - "per_seed": legacy path — N independent (build → transform → enc → # den) passes, sequential because hulls differ. Used as fallback when # the union-hull build fails (no valid hull). seed_results: List[Tuple[np.ndarray, List[Tuple[int, int]], np.ndarray]] = [] n_hull_valid = 0 if strategy == "union_hull" and n_ensemble > 1: union_results, hull_valid = _predict_union_hull_ensemble( raw_sample=raw_sample, colmap_rec=colmap_rec, s1_runs=s1_runs, stage2_cfg=stage2_cfg, device=device, preprocess_cache=preprocess_cache, timings_accum=timings, order_id=order_id, ) if hull_valid and union_results: n_hull_valid = n_ensemble for i, (v, e, l) in enumerate(union_results): if len(v) > 0 and len(e) > 0: seed_results.append(( np.asarray(v, dtype=float), list(e), np.asarray(l, dtype=np.float32), )) else: # This slot's stage-2 returned nothing — fall back to its # stage-1 prediction so it still contributes to selection. s1_v, s1_e_arr, s1_l = s1_runs[i] s1_e_list = [(int(a), int(b)) for a, b in s1_e_arr.tolist()] if len(s1_v) > 0 and len(s1_e_list) > 0: seed_results.append(( np.asarray(s1_v, dtype=float), s1_e_list, np.asarray(s1_l, dtype=np.float32), )) # If hull was invalid OR union path produced nothing, fall through # to per-seed below so we don't silently lose this sample. if not seed_results: for s1_verts, s1_edges_arr, s1_logits in s1_runs: s2_verts, s2_edges, s2_logits, hull_valid = _run_stage2_for_seed( raw_sample=raw_sample, colmap_rec=colmap_rec, pred_verts=s1_verts, pred_edges=s1_edges_arr, stage2_cfg=stage2_cfg, device=device, preprocess_cache=preprocess_cache, timings_accum=timings, order_id=order_id, ) if hull_valid: n_hull_valid += 1 if (s2_verts is not None and s2_edges is not None and len(s2_verts) > 0 and len(s2_edges) > 0): seed_results.append(( np.asarray(s2_verts, dtype=float), list(s2_edges), np.asarray(s2_logits, dtype=np.float32), )) else: s1_e_list = [(int(a), int(b)) for a, b in s1_edges_arr.tolist()] if len(s1_verts) > 0 and len(s1_e_list) > 0: seed_results.append(( np.asarray(s1_verts, dtype=float), s1_e_list, np.asarray(s1_logits, dtype=np.float32), )) # 3) Fuse / select. When n_ensemble=1 the single run is returned verbatim. if not seed_results: if stats is not None: stats["empty_fallback"] = stats.get("empty_fallback", 0) + 1 if diagnostics is not None: diagnostics["members"] = [] diagnostics["final"] = empty_solution() ev, ee = empty_solution() return np.asarray(ev, dtype=float), [(int(a), int(b)) for a, b in ee] if diagnostics is not None: diagnostics["members"] = [ (np.asarray(v, dtype=float), list(e), np.asarray(l, dtype=np.float32)) for v, e, l in seed_results ] diagnostics["s1_runs"] = [ (np.asarray(v, dtype=float), list(e), np.asarray(l, dtype=np.float32)) for v, e, l in s1_runs ] diagnostics["n_hull_valid"] = int(n_hull_valid) if n_ensemble == 1: fused_v, fused_e = seed_results[0][0], seed_results[0][1] if diagnostics is not None: diagnostics["single"] = (np.asarray(fused_v, dtype=float), list(fused_e)) else: selector = str(getattr(stage2_cfg, "ensemble_selector", "legacy")) mode = str(getattr(stage2_cfg, "ensemble_mode", "medoid")) tau = float(getattr(stage2_cfg, "ensemble_merge_m", 0.4)) refine = bool(getattr(stage2_cfg, "ensemble_refine_positions", True)) top_k = int(getattr(stage2_cfg, "ensemble_top_k", 0)) edge_vote = bool(getattr(stage2_cfg, "ensemble_edge_vote", False)) edge_vote_frac = float(getattr(stage2_cfg, "ensemble_edge_vote_frac", 0.5)) vertex_vote_frac = float(getattr(stage2_cfg, "ensemble_vertex_vote_frac", 0.0)) agg = str(getattr(stage2_cfg, "ensemble_refine_agg", "mean")) ranking: Optional[List[int]] = None if mode == "medoid" or top_k > 0 or selector == "ranker": ranking = _medoid_rank(seed_results, match_tau_m=tau) medoid_pick = (ranking[0] if ranking else _medoid_select_idx(seed_results, match_tau_m=tau)) conf_pick = _confidence_select_idx(seed_results) if diagnostics is not None: diagnostics["single"] = ( np.asarray(seed_results[0][0], dtype=float), list(seed_results[0][1]), ) diagnostics["medoid_pick"] = ( np.asarray(seed_results[medoid_pick][0], dtype=float), list(seed_results[medoid_pick][1]), ) diagnostics["confidence_pick"] = ( np.asarray(seed_results[conf_pick][0], dtype=float), list(seed_results[conf_pick][1]), ) diagnostics["medoid_idx"] = int(medoid_pick) diagnostics["confidence_idx"] = int(conf_pick) if selector in ("ranker", "fixed"): reproj_views_cached = _build_reprojection_views( raw_sample, colmap_rec, ) if selector == "ranker" else None candidates = _build_ranker_candidates( seed_results, match_tau_m=tau, edge_vote_frac=edge_vote_frac, ranking=ranking, medoid_idx=medoid_pick, confidence_idx=conf_pick, reproj_views=reproj_views_cached, s1_runs=s1_runs, ) if candidates: if selector == "fixed": wanted = str(getattr( stage2_cfg, "ensemble_fixed_candidate", "confidence_mean_add")) fallback = "medoid_mean_add" scores = [ (2.0 if c["name"] == wanted else 1.0 if c["name"] == fallback else 0.0) for c in candidates ] else: ranker = getattr(stage2_cfg, "ensemble_ranker", None) scores = [ _score_ranker_candidate(c["features"], ranker) for c in candidates ] scores_arr = np.asarray(scores, dtype=np.float64) best_c = int(np.argmax(scores_arr)) chosen = candidates[best_c] topk_k = int(getattr(stage2_cfg, "ensemble_topk_fuse", 1)) if (selector == "ranker" and topk_k > 1 and len(candidates) > 1): fused_v, fused_e, _ = _fuse_topk_candidates( candidates, scores_arr, topk_k, match_tau_m=tau, edge_vote_frac=edge_vote_frac, ) else: fused_v = np.asarray(chosen["verts"], dtype=np.float64) fused_e = list(chosen["edges"]) if diagnostics is not None: diagnostics["candidates"] = [ { "name": str(c["name"]), "features": np.asarray(c["features"], dtype=np.float64), "score": float(scores[j]), "verts": np.asarray(c["verts"], dtype=float), "edges": list(c["edges"]), } for j, c in enumerate(candidates) ] diagnostics["selected_candidate"] = str(chosen["name"]) diagnostics["selected_candidate_score"] = float(scores[best_c]) diagnostics["pick_idx"] = int(chosen.get("source_idx", -1)) diagnostics["topk_fuse_k"] = int(topk_k) else: fused_v = np.asarray(seed_results[medoid_pick][0], dtype=np.float64) fused_e = list(seed_results[medoid_pick][1]) if diagnostics is not None: diagnostics["pick_idx"] = int(medoid_pick) elif mode == "consensus": # Consensus already aggregates across runs — refinement is # logically subsumed. Vertex / edge vote fractions come straight # from the same flags (vertex frac defaults to 0.5 here since # consensus needs *some* support bar, not 0). v_frac_consensus = vertex_vote_frac if vertex_vote_frac > 0.0 else 0.5 e_frac_consensus = edge_vote_frac if edge_vote else 0.5 fused_v, fused_e = _fuse_wireframes( list(seed_results), n_total=n_ensemble, merge_tau_m=tau, vertex_vote_frac=v_frac_consensus, edge_vote_frac=e_frac_consensus, aggregator=agg, ) else: if mode == "confidence": pick = conf_pick else: # "medoid" pick = medoid_pick if diagnostics is not None: diagnostics["pick_idx"] = int(pick) if top_k > 0 and ranking: members: Optional[List[int]] = list(ranking[:max(1, top_k)]) if pick not in members: members.append(pick) else: members = None # use all runs keep_idx: Optional[np.ndarray] = None if refine: fused_v, fused_e, keep_idx = _refine_positions( seed_results, pick, match_tau_m=tau, aggregator=agg, member_indices=members, vertex_vote_frac=vertex_vote_frac, ) else: fused_v = np.asarray(seed_results[pick][0], dtype=np.float64).reshape(-1, 3) fused_e = list(seed_results[pick][1]) keep_idx = np.arange(fused_v.shape[0], dtype=np.int64) if edge_vote: fused_e = _vote_edges( seed_results, pick, match_tau_m=tau, vote_frac=edge_vote_frac, member_indices=members, keep_idx=keep_idx, ) if stats is not None: if n_hull_valid == 0: stats["stage1_fallback"] = stats.get("stage1_fallback", 0) + 1 else: stats["stage2_used"] = stats.get("stage2_used", 0) + 1 if len(fused_v) == 0 or len(fused_e) == 0: if stats is not None: stats["empty_fallback"] = stats.get("empty_fallback", 0) + 1 if diagnostics is not None: diagnostics["final"] = empty_solution() ev, ee = empty_solution() return np.asarray(ev, dtype=float), [(int(a), int(b)) for a, b in ee] if diagnostics is not None: diagnostics["final"] = (np.asarray(fused_v, dtype=float), list(fused_e)) return np.asarray(fused_v, dtype=float), fused_e # --------------------------------------------------------------------------- # Dataset loading (mirrors empty_submission_2026) # --------------------------------------------------------------------------- def load_test_dataset(params: Dict[str, Any]): if load_dataset is None: raise SystemExit("[data] missing Python package 'datasets'") data_path_test_server = Path("/tmp/data") if data_path_test_server.exists(): print("[data] test environment detected (/tmp/data exists)", flush=True) else: print(f"[data] local run — snapshot_download '{params['dataset']}' → /tmp/data", flush=True) from huggingface_hub import snapshot_download snapshot_download( repo_id=params["dataset"], local_dir="/tmp/data", repo_type="dataset", ) data_path = data_path_test_server data_files = { "validation": [str(p) for p in data_path.rglob("*public*/**/*.tar")], "test": [str(p) for p in data_path.rglob("*private*/**/*.tar")], } print(f"[data] resolved data_files: " f"{ {k: len(v) for k, v in data_files.items()} } shards", flush=True) return load_dataset( str(data_path / "hoho22k_2026_test_x_anon.py"), data_files=data_files, trust_remote_code=True, writer_batch_size=100, ) # --------------------------------------------------------------------------- # Inference loop with prefetched preprocessing # --------------------------------------------------------------------------- def predict_split( split_ds, model: WireframeDiffusion, device: torch.device, n_pts: int, use_depth: bool, n_sample_steps: int, validity_thresh: float, n_prefetch: int, stage2_cfg: Stage2Config, ) -> List[Dict[str, Any]]: """Sequential GPU inference; CPU preprocessing prefetched n_prefetch ahead. Stage 1 runs on the V10 scene, its prediction defines the hull, the stage-2 scene is built from the hull-cropped raw points, and stage 2 produces the final wireframe. The CPU prefetch holds onto the raw sample + decoded COLMAP record so the stage-2 builder doesn't re-parse the colmap zip. """ results: List[Dict[str, Any]] = [] use_stage2 = stage2_cfg is not None pool = ThreadPoolExecutor(max_workers=max(1, n_prefetch)) pending: "queue.Queue[Tuple[Any, Any]]" = queue.Queue() sample_iter = iter(split_ds) def submit_next() -> bool: try: sample = next(sample_iter) except StopIteration: return False order_id = sample.get("order_id") fut = pool.submit(preprocess_sample, sample, n_pts, use_depth, 0.1, use_stage2) pending.put((order_id, fut)) return True for _ in range(max(1, n_prefetch)): if not submit_next(): break pbar = tqdm(desc="predict", unit="sample") while not pending.empty(): order_id, fut = pending.get() submit_next() t0 = time.perf_counter() try: pre = fut.result() if use_stage2: scene = pre["scene"] raw_sample = pre["raw_sample"] colmap_rec = pre["colmap_rec"] cache = pre.get("cache") verts, edges = predict_two_stage( raw_sample=raw_sample, colmap_rec=colmap_rec, stage1_scene=scene, stage1_model=model, stage1_n_steps=n_sample_steps, stage1_validity_thresh=validity_thresh, stage2_cfg=stage2_cfg, device=device, preprocess_cache=cache, ) verts = np.asarray(verts, dtype=float) else: scene = pre batch = scene_to_batch(scene, device) with torch.inference_mode(): verts, edges = model.predict_wireframe( batch, n_steps=n_sample_steps, validity_thresh=validity_thresh, ) verts = np.asarray(verts, dtype=float) edges = [(int(a), int(b)) for a, b in edges] dt = time.perf_counter() - t0 tag = "stage2" if use_stage2 else "stage1" print(f"[predict/{tag}] {order_id}: {len(verts)}v {len(edges)}e ({dt:.2f}s)", flush=True) except Exception as exc: print(f"[predict] {order_id}: FAILED ({type(exc).__name__}: {exc}) " f"— empty_solution", flush=True) verts, edges = empty_solution() edges = [(int(a), int(b)) for a, b in edges] results.append({ "order_id": order_id, "wf_vertices": verts.tolist(), "wf_edges": edges, }) pbar.update(1) pbar.close() pool.shutdown(wait=True) return results # --------------------------------------------------------------------------- # Local sanity check (--sanity): val split, scored against GT # --------------------------------------------------------------------------- def run_sanity_check( model: WireframeDiffusion, device: torch.device, n_pts: int, use_depth: bool, n_sample_steps: int, validity_thresh: float, n_samples: int, split: str, dataset_id: str, cache_dir: str, stage2_cfg: Stage2Config, ensemble_diagnostics: bool = False, candidate_dump_path: Optional[str] = None, ) -> float: """Run the model on a handful of val samples with GT and report HSS. Streams from `dataset_id` (default the public hoho22k trainval split) so we never need /tmp/data or the anonymised test shards. Returns mean HSS across the evaluated samples. """ import numpy as np from hoho2025.metric_helper import hss as hss_fn, HSSReturnType use_all = n_samples <= 0 n_label = "all" if use_all else str(n_samples) print(f"[sanity] dataset={dataset_id} split={split} n={n_label} " f"n_pts={n_pts} use_depth={use_depth} " f"n_sample_steps={n_sample_steps} validity_thresh={validity_thresh}", flush=True) t_load = time.perf_counter() if load_dataset is None: raise SystemExit("[sanity] missing Python package 'datasets'") ds = load_dataset( dataset_id, streaming=True, trust_remote_code=True, cache_dir=str(Path(cache_dir) / "hf"), ) print(f"[sanity] dataset ready in {time.perf_counter() - t_load:.1f}s " f"(available splits: {list(ds.keys())})", flush=True) if split not in ds: raise SystemExit(f"[sanity] split '{split}' not in dataset; " f"available: {list(ds.keys())}") split_ds = ds[split] hss_scores: List[float] = [] f1_scores: List[float] = [] iou_scores: List[float] = [] diag_hss: Dict[str, List[float]] = {} diag_counts: Dict[str, List[int]] = {} candidate_dump_f = open(candidate_dump_path, "w") if candidate_dump_path else None if candidate_dump_f is not None: print(f"[ranker/dump] writing candidates to {candidate_dump_path}", flush=True) t_pre_list: List[float] = [] t_enc_list: List[float] = [] t_den_list: List[float] = [] def _sync(): if device.type == "cuda": torch.cuda.synchronize() use_stage2 = stage2_cfg is not None if use_stage2: print(f"[sanity] STAGE 2 enabled n_pts={stage2_cfg.n_pts} " f"n_sample_steps={stage2_cfg.n_sample_steps} " f"hull=(+{stage2_cfg.above_m:.2f}/-{stage2_cfg.below_m:.2f}/" f"±{stage2_cfg.side_m:.2f})", flush=True) if ensemble_diagnostics and stage2_cfg.ensemble_n > 1: print("[diag] enabled: scoring single/member0, medoid pick, " "confidence pick, final fused, and oracle-best member", flush=True) t_s2_list: List[float] = [] t_s2_build_list: List[float] = [] t_s2_xform_list: List[float] = [] t_s2_enc_list: List[float] = [] t_s2_den_list: List[float] = [] stage2_stats: Dict[str, int] = {} for i, sample in enumerate(split_ds): if not use_all and i >= n_samples: break order_id = sample.get("order_id", i) t_pre = t_enc = t_den = t_s2 = float("nan") s2_timings: Dict[str, float] = {} t_total0 = time.perf_counter() try: diag: Dict[str, Any] = {} # 1) Preprocessing (CPU) + H2D transfer t0 = time.perf_counter() pre = preprocess_sample(sample, n_pts=n_pts, use_depth=use_depth, keep_raw=use_stage2) scene = pre["scene"] if use_stage2 else pre batch = scene_to_batch(scene, device) _sync() t_pre = time.perf_counter() - t0 with torch.inference_mode(): # 2) Scene encoder t0 = time.perf_counter() scene_feats, scene_xyz_norm = model._encode_scene(batch) _sync() t_enc = time.perf_counter() - t0 # 3) Vertex denoiser — n_sample_steps Euler steps K = model.denoiser.k_verts B = scene_feats.shape[0] step_dt = 1.0 / n_sample_steps x = model._init_x0(batch, K, device, B) last_logit = torch.zeros(B, K, device=device) last_edge_logit = torch.zeros(B, K, K, device=device) t0 = time.perf_counter() for s in range(n_sample_steps): t = torch.full((B,), s * step_dt, device=device) v, last_logit, last_edge_logit = model.denoiser( x, t, scene_feats, scene_xyz_norm, ) x = x + step_dt * v _sync() t_den = time.perf_counter() - t0 # 4) Postprocess → world-space verts/edges (mirrors # WireframeDiffusion.predict_wireframe; not timed). xyz_norm = x[0] valid = (last_logit[0] > validity_thresh) edge_logit = last_edge_logit[0].float() valid_idx = torch.nonzero(valid, as_tuple=False).flatten() xyz_valid = xyz_norm.index_select(0, valid_idx) center = batch["bbox_center"][0].to(xyz_valid) scale = batch["bbox_scale"][0].to(xyz_valid) bbox_R = batch.get("bbox_R") if isinstance(bbox_R, torch.Tensor): R = bbox_R[0].to(xyz_valid) verts_world = (xyz_valid * scale) @ R + center else: verts_world = xyz_valid * scale + center pred_v = verts_world.cpu().numpy() pred_e: List[Tuple[int, int]] = [] n_valid = int(valid_idx.numel()) if n_valid >= 2: sub = edge_logit.index_select(0, valid_idx).index_select(1, valid_idx) tri = torch.triu(torch.ones(n_valid, n_valid, device=sub.device, dtype=torch.bool), diagonal=1) pairs = torch.nonzero((sub > 0.0) & tri, as_tuple=False) pred_e = [(int(a), int(b)) for a, b in pairs.detach().cpu().tolist()] # 5) Stage 2 (optional) — build cropped scene + refine. The # `s2_timings` dict gets populated with per-phase wallclock # so we can print stage-2 with the same breakdown as stage 1. if use_stage2: t0 = time.perf_counter() pred_v_s2, pred_e_s2 = predict_two_stage( raw_sample=pre["raw_sample"], colmap_rec=pre["colmap_rec"], stage1_scene=scene, stage1_model=model, stage1_n_steps=n_sample_steps, stage1_validity_thresh=validity_thresh, stage2_cfg=stage2_cfg, device=device, preprocess_cache=pre.get("cache"), stats=stage2_stats, timings=s2_timings, diagnostics=(diag if ensemble_diagnostics else None), ) _sync() t_s2 = time.perf_counter() - t0 pred_v = np.asarray(pred_v_s2, dtype=float) pred_e = list(pred_e_s2) except Exception as exc: print(f"[sanity] {order_id}: PREDICT FAILED " f"({type(exc).__name__}: {exc})", flush=True) pred_v, pred_e = empty_solution() gt_v = np.asarray(sample["wf_vertices"]) gt_e = np.asarray(sample["wf_edges"]) def _score_diag_variant(name: str, verts: np.ndarray, edges: List[Tuple[int, int]]) -> float: try: r = hss_fn(verts, edges, gt_v, gt_e) val = float(r.hss) except Exception: val = 0.0 diag_hss.setdefault(name, []).append(val) diag_counts.setdefault(f"{name}_v", []).append(int(len(verts))) diag_counts.setdefault(f"{name}_e", []).append(int(len(edges))) return val diag_line = "" if use_stage2 and ensemble_diagnostics and "diag" in locals() and diag: scored: Dict[str, float] = {} for name in ("single", "medoid_pick", "confidence_pick", "final"): if name in diag: vv, ee = diag[name] scored[name] = _score_diag_variant( name, np.asarray(vv, dtype=float), list(ee)) member_scores: List[float] = [] for vv, ee, _ll in diag.get("members", []): member_scores.append(_score_diag_variant( "member", np.asarray(vv, dtype=float), list(ee))) if member_scores: oracle = max(member_scores) diag_hss.setdefault("oracle", []).append(float(oracle)) best_idx = int(np.argmax(np.asarray(member_scores, dtype=np.float64))) final_h = scored.get("final", 0.0) med_h = scored.get("medoid_pick", 0.0) conf_h = scored.get("confidence_pick", 0.0) single_h = scored.get("single", 0.0) diag_line = ( f" | diag single={single_h:.3f} medoid={med_h:.3f} " f"conf={conf_h:.3f} final={final_h:.3f} " f"oracle={oracle:.3f}@{best_idx} gap={oracle-final_h:+.3f}" ) if candidate_dump_f is not None and diag.get("members"): candidates = diag.get("candidates") if not candidates: members_for_dump = [ (np.asarray(vv, dtype=float), list(ee), np.asarray(ll, dtype=np.float32)) for vv, ee, ll in diag.get("members", []) ] s1_for_dump = [ (np.asarray(vv, dtype=float), list(ee), np.asarray(ll, dtype=np.float32)) for vv, ee, ll in diag.get("s1_runs", []) ] candidates = _build_ranker_candidates( members_for_dump, match_tau_m=stage2_cfg.ensemble_merge_m, edge_vote_frac=stage2_cfg.ensemble_edge_vote_frac, medoid_idx=diag.get("medoid_idx"), confidence_idx=diag.get("confidence_idx"), reproj_views=_build_reprojection_views( pre["raw_sample"], pre["colmap_rec"], ), s1_runs=s1_for_dump, ) selected_name = str(diag.get("selected_candidate", "")) for cand in candidates: cv = np.asarray(cand["verts"], dtype=float) ce = list(cand["edges"]) try: cr = hss_fn(cv, ce, gt_v, gt_e) chss = float(cr.hss) cf1 = float(cr.f1) ciou = float(cr.iou) except Exception: chss, cf1, ciou = 0.0, 0.0, 0.0 row = { "order_id": order_id, "sample_index": int(i), "name": str(cand["name"]), "selected": bool(str(cand["name"]) == selected_name), "hss": chss, "f1": cf1, "iou": ciou, "n_v": int(len(cv)), "n_e": int(len(ce)), "features": np.asarray( cand["features"], dtype=np.float64).tolist(), } candidate_dump_f.write(json.dumps(row, separators=(",", ":")) + "\n") candidate_dump_f.flush() try: result = hss_fn(pred_v, pred_e, gt_v, gt_e) except Exception as exc: print(f"[sanity] {order_id}: HSS FAILED " f"({type(exc).__name__}: {exc})", flush=True) result = HSSReturnType(hss=0.0, f1=0.0, iou=0.0) hss_scores.append(float(result.hss)) f1_scores.append(float(result.f1)) iou_scores.append(float(result.iou)) if t_pre == t_pre: t_pre_list.append(t_pre) # filter NaNs if t_enc == t_enc: t_enc_list.append(t_enc) if t_den == t_den: t_den_list.append(t_den) if t_s2 == t_s2: t_s2_list.append(t_s2) if use_stage2: if "s2_build" in s2_timings: t_s2_build_list.append(s2_timings["s2_build"]) if "s2_transform" in s2_timings: t_s2_xform_list.append(s2_timings["s2_transform"]) if "s2_enc" in s2_timings: t_s2_enc_list.append(s2_timings["s2_enc"]) if "s2_den" in s2_timings: t_s2_den_list.append(s2_timings["s2_den"]) dt = time.perf_counter() - t_total0 if use_stage2 and t_s2 == t_s2: s2_steps = stage2_cfg.n_sample_steps s2_str = (f" s2={t_s2:.2f}s " f"(build={s2_timings.get('s2_build', float('nan')):.2f}s " f"xform={s2_timings.get('s2_transform', float('nan')):.3f}s " f"enc={s2_timings.get('s2_enc', float('nan')):.3f}s " f"den[{s2_steps}]={s2_timings.get('s2_den', float('nan')):.3f}s)") else: s2_str = "" print(f"[sanity] [{i:03d}] {order_id} pred={len(pred_v)}v/{len(pred_e)}e " f"gt={len(gt_v)}v/{len(gt_e)}e " f"HSS={result.hss:.3f} F1={result.f1:.3f} IoU={result.iou:.3f} " f"| pre={t_pre:.2f}s enc={t_enc:.3f}s " f"den[{n_sample_steps}]={t_den:.3f}s{s2_str} total={dt:.2f}s" f"{diag_line}", flush=True) def _mean_std(xs: List[float]) -> Tuple[float, float]: if not xs: return 0.0, 0.0 arr = np.asarray(xs, dtype=np.float64) return float(arr.mean()), float(arr.std()) n = len(hss_scores) hss_m, hss_s = _mean_std(hss_scores) f1_m, f1_s = _mean_std(f1_scores) iou_m, iou_s = _mean_std(iou_scores) print("=" * 60, flush=True) print(f"[sanity] samples={n}", flush=True) print(f"[sanity] HSS mean={hss_m:.4f} std={hss_s:.4f}", flush=True) print(f"[sanity] F1 mean={f1_m:.4f} std={f1_s:.4f}", flush=True) print(f"[sanity] IoU mean={iou_m:.4f} std={iou_s:.4f}", flush=True) if diag_hss: print("[diag] HSS means:", flush=True) for name in ("single", "member", "medoid_pick", "confidence_pick", "final", "oracle"): if name not in diag_hss: continue m, s = _mean_std(diag_hss[name]) print(f"[diag] {name:15s} mean={m:.4f} std={s:.4f}", flush=True) if "oracle" in diag_hss and "final" in diag_hss: gap = (np.asarray(diag_hss["oracle"], dtype=np.float64) - np.asarray(diag_hss["final"], dtype=np.float64)) print(f"[diag] oracle-final mean={float(gap.mean()):+.4f} " f"std={float(gap.std()):.4f}", flush=True) for name in ("single", "medoid_pick", "confidence_pick", "final"): vk, ek = f"{name}_v", f"{name}_e" if vk in diag_counts and ek in diag_counts: v_m, v_s = _mean_std([float(x) for x in diag_counts[vk]]) e_m, e_s = _mean_std([float(x) for x in diag_counts[ek]]) print(f"[diag] {name:15s} size={v_m:.1f}±{v_s:.1f}v/" f"{e_m:.1f}±{e_s:.1f}e", flush=True) if use_all: print(f"[sanity] full split evaluated: {n} samples", flush=True) if t_pre_list: pre_m, pre_s = _mean_std(t_pre_list) enc_m, enc_s = _mean_std(t_enc_list) den_m, den_s = _mean_std(t_den_list) per_step_m = (den_m / n_sample_steps) if n_sample_steps else 0.0 per_step_s = (den_s / n_sample_steps) if n_sample_steps else 0.0 print(f"[sanity] preprocess mean={pre_m:.3f}s std={pre_s:.3f}s", flush=True) print(f"[sanity] encoder mean={enc_m:.3f}s std={enc_s:.3f}s", flush=True) print(f"[sanity] denoiser[{n_sample_steps}] mean={den_m:.3f}s std={den_s:.3f}s " f"({per_step_m*1000:.1f}±{per_step_s*1000:.1f} ms/step)", flush=True) if use_stage2 and t_s2_list: s2_m, s2_s = _mean_std(t_s2_list) print(f"[sanity] stage-2 total mean={s2_m:.3f}s std={s2_s:.3f}s", flush=True) if t_s2_build_list: b_m, b_s = _mean_std(t_s2_build_list) print(f"[sanity] stage-2 build mean={b_m:.3f}s std={b_s:.3f}s", flush=True) if t_s2_xform_list: x_m, x_s = _mean_std(t_s2_xform_list) print(f"[sanity] stage-2 transform mean={x_m:.3f}s std={x_s:.3f}s", flush=True) if t_s2_enc_list: e_m, e_s = _mean_std(t_s2_enc_list) print(f"[sanity] stage-2 encoder mean={e_m:.3f}s std={e_s:.3f}s", flush=True) if t_s2_den_list: d_m, d_s = _mean_std(t_s2_den_list) s2_steps = stage2_cfg.n_sample_steps per_m = (d_m / s2_steps) if s2_steps else 0.0 per_s = (d_s / s2_steps) if s2_steps else 0.0 print(f"[sanity] stage-2 denoiser[{s2_steps}] mean={d_m:.3f}s std={d_s:.3f}s " f"({per_m*1000:.1f}±{per_s*1000:.1f} ms/step)", flush=True) print(f"[sanity] stage-2 used={stage2_stats.get('stage2_used', 0)} " f"stage-1 fallback={stage2_stats.get('stage1_fallback', 0)} " f"empty fallback={stage2_stats.get('empty_fallback', 0)}", flush=True) if candidate_dump_f is not None: candidate_dump_f.close() print(f"[ranker/dump] closed {candidate_dump_path}", flush=True) print("=" * 60, flush=True) return hss_m # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def _build_stage2_cfg(args, params: Dict[str, Any], device: torch.device) -> Stage2Config: """Resolve --stage2_ckpt / env / params, load the Stage2Diffusion model, and assemble a ``Stage2Config``. The two-stage pipeline is mandatory; ``resolve_stage2_ckpt_path`` raises if no ckpt is configured.""" ckpt_path = resolve_stage2_ckpt_path(args.stage2_ckpt, params) model, s2_args = load_stage2_model(ckpt_path, device) n_pts = (args.stage2_n_pts if args.stage2_n_pts is not None else s2_args.get("n_pts", 16384)) k_verts = s2_args.get("k_verts", 64) n_sample_steps = (args.stage2_n_sample_steps if args.stage2_n_sample_steps is not None else s2_args.get("n_sample_steps", 20)) validity_thresh = (args.stage2_validity_thresh if args.stage2_validity_thresh is not None else s2_args.get("validity_thresh", 0.0)) # Keep stage-2 inference on the original random sampler. Earlier FPS # experiments added extra plumbing here, but the current submission/eval # path intentionally does not switch sampling modes from CLI or ckpt args. sampling = "random" # Oversample budgets: random uses a smaller candidate pool; voxel is left # as a code path for old experiments but is not selected by this script. def _default_oversample(s: str) -> int: if s == "voxel": return 16_000 return 10_000 n_col_oversample = (args.stage2_n_col_oversample if args.stage2_n_col_oversample is not None else _default_oversample(sampling)) n_dep_oversample = (args.stage2_n_dep_oversample if args.stage2_n_dep_oversample is not None else _default_oversample(sampling)) fps_max_exact_iters = None cfg = Stage2Config( model=model, n_pts=n_pts, k_verts=k_verts, n_sample_steps=n_sample_steps, validity_thresh=validity_thresh, above_m=args.stage2_above_m, below_m=args.stage2_below_m, side_m=args.stage2_side_m, n_col_oversample=n_col_oversample, n_dep_oversample=n_dep_oversample, n_depth_per_image_cap=args.stage2_n_depth_per_image_cap, sampling=sampling, density_voxel_size_m=args.stage2_density_voxel_size_m, density_kernel_radius=args.stage2_density_kernel_radius, density_kernel_axis=args.stage2_density_kernel_axis, density_response_power=args.stage2_density_response_power, density_planarity_suppression=args.stage2_density_planarity_suppression, density_planarity_radius=args.stage2_density_planarity_radius, density_planarity_min_points=args.stage2_density_planarity_min_points, density_min_per_voxel=args.stage2_density_min_per_voxel, ensemble_n=args.ensemble_n, ensemble_merge_m=args.ensemble_merge_m, ensemble_mode=args.ensemble_mode, ensemble_strategy=args.ensemble_strategy, ensemble_refine_positions=(not args.no_position_refine), ensemble_refine_agg=args.ensemble_refine_agg, ensemble_top_k=args.ensemble_top_k, ensemble_edge_vote=(not args.no_edge_vote), ensemble_edge_vote_frac=args.ensemble_edge_vote_frac, ensemble_vertex_vote_frac=args.ensemble_vertex_vote_frac, ensemble_stage2_last_k=args.ensemble_stage2_last_k, ensemble_selector=args.ensemble_selector, ensemble_ranker_path=args.ensemble_ranker_path, ensemble_ranker=( _load_ensemble_ranker(args.ensemble_ranker_path) if args.ensemble_selector == "ranker" else None ), ensemble_fixed_candidate=args.ensemble_fixed_candidate, ensemble_topk_fuse=args.ensemble_topk_fuse, fps_max_exact_iters=fps_max_exact_iters, ) extra = "" if cfg.sampling == "voxel": extra = (f" density(voxel={cfg.density_voxel_size_m:.2f}m " f"kernel={cfg.density_kernel_axis}:r{cfg.density_kernel_radius} " f"planarity_supp={cfg.density_planarity_suppression:.2f})") elif cfg.sampling == "fps": extra = (f" fps_max_exact_iters=" f"{'full' if cfg.fps_max_exact_iters is None else cfg.fps_max_exact_iters}") print(f"[stage2] enabled n_pts={cfg.n_pts} steps={cfg.n_sample_steps} " f"sampling={cfg.sampling} oversample={cfg.n_col_oversample}+{cfg.n_dep_oversample} " f"hull=(+{cfg.above_m:.2f}/-{cfg.below_m:.2f}/±{cfg.side_m:.2f}){extra}", flush=True) if cfg.ensemble_n > 1: refine_str = (cfg.ensemble_refine_agg if cfg.ensemble_refine_positions else "off") topk_str = ("all" if cfg.ensemble_top_k == 0 else f"top{cfg.ensemble_top_k}") edge_str = (f"vote>={cfg.ensemble_edge_vote_frac:.2f}" if cfg.ensemble_edge_vote else "picked") vdrop_str = (f"vote>={cfg.ensemble_vertex_vote_frac:.2f}" if cfg.ensemble_vertex_vote_frac > 0.0 else "off") print(f"[ensemble] N={cfg.ensemble_n} strategy={cfg.ensemble_strategy} " f"mode={cfg.ensemble_mode} tau={cfg.ensemble_merge_m:.2f}m " f"refine={refine_str} members={topk_str} edges={edge_str} " f"vdrop={vdrop_str} s2_lastk={cfg.ensemble_stage2_last_k} " f"selector={cfg.ensemble_selector}", flush=True) return cfg def parse_args(argv=None): p = argparse.ArgumentParser(description="S23DR 2026 submission script (trained model)") p.add_argument("--ckpt", default=None, help="Checkpoint path; falls back to $S23DR_CKPT, params['ckpt'], " "then ./test_checkpoint.pth") p.add_argument("--params", default="params.json") p.add_argument("--n_sample_steps", type=int, default=50, help="Stage-1 diffusion ODE steps at inference (default: 50).") p.add_argument("--n_pts", type=int, default=None, help="Stage-1 scene points at inference (default: ckpt args / 8192)") p.add_argument("--validity_thresh", type=float, default=0.0, help="Drop predicted vertices below this validity logit " "(default: 0.0)") p.add_argument("--n_prefetch", type=int, default=2, help="Number of samples to preprocess ahead of the GPU (default: 2)") p.add_argument("--device", default="cuda") p.add_argument("--output", default="submission.json") p.add_argument("--limit", type=int, default=0, help="If >0, only process the first N samples per split (debug)") # ---- stage-2 refinement (optional) ----------------------------------- p.add_argument("--stage2_ckpt", default=None, help="Stage-2 refinement checkpoint. Falls back to " "$S23DR_STAGE2_CKPT then params['stage2_ckpt']. " "When unset, runs stage-1 only.") p.add_argument("--stage2_n_pts", type=int, default=None, help="Stage-2 scene points (default: ckpt args / 16384).") p.add_argument("--stage2_n_sample_steps", type=int, default=50, help="Stage-2 ODE steps (default: 50 — 2.5x training, " "inference oversampling).") p.add_argument("--stage2_validity_thresh", type=float, default=0.0, help="Stage-2 validity logit threshold (default: 0.0).") p.add_argument("--stage2_above_m", type=float, default=0.5, help="Hull inflation upward (metres). Default 0.5.") p.add_argument("--stage2_below_m", type=float, default=0.5, help="Hull inflation downward (metres). Default 0.5.") p.add_argument("--stage2_side_m", type=float, default=0.5, help="Hull inflation laterally in X/Y (metres). Default 0.5.") p.add_argument("--stage2_n_col_oversample", type=int, default=None, help="Stage-2 COLMAP oversample before sampling drops to n_pts//2. " "Default: 10k for random sampling.") p.add_argument("--stage2_n_dep_oversample", type=int, default=None, help="Stage-2 depth oversample before sampling drops to n_pts//2. " "Default: 10k for random sampling.") p.add_argument("--stage2_n_depth_per_image_cap", type=int, default=30_000, help="Cap on per-image depth pixels considered before the hull test.") p.add_argument("--stage2_sampling", choices=["random"], default="random", help="Stage-2 scene sampling mode. The eval/submission path " "is locked to random sampling.") p.add_argument("--stage2_fps_max_exact_iters", type=int, default=None, help=argparse.SUPPRESS) # ---- stage-2 density-contrast hyperparams (only used when --stage2_sampling=voxel) - p.add_argument("--stage2_density_voxel_size_m", type=float, default=0.25) p.add_argument("--stage2_density_kernel_radius", type=int, default=3) p.add_argument("--stage2_density_kernel_axis", choices=["cube", "x", "y", "z"], default="cube") p.add_argument("--stage2_density_response_power", type=float, default=1.0) p.add_argument("--stage2_density_planarity_suppression", type=float, default=1.0) p.add_argument("--stage2_density_planarity_radius", type=int, default=1) p.add_argument("--stage2_density_planarity_min_points", type=int, default=12) p.add_argument("--stage2_density_min_per_voxel", type=int, default=0) # ---- inference-time ensembling --------------------------------------- p.add_argument("--ensemble_n", type=int, default=16, help="Run N stage-1 random-init trajectories and combine " "them per --ensemble_mode. Default 16 (paired with " "--ensemble_strategy union_hull this is ~1.2x wall " "since stage-2 shares one encoder pass at B=N). " "Set to 1 to disable.") p.add_argument("--ensemble_merge_m", type=float, default=0.5, help="Vertex match threshold in metres for ensemble " "selection (medoid Hungarian distance, position " "refinement correspondents, consensus clustering). " "Default 0.5 — matches the HSS metric's vert_thresh.") p.add_argument("--ensemble_strategy", choices=["union_hull", "per_seed"], default="union_hull", help="How stage-2 is run across the N ensemble members. " "'union_hull' (default) builds ONE stage-2 scene " "from the union of all stage-1 vertex predictions, " "encodes once, and runs a B=N denoiser batch with " "per-slot init_verts. 'per_seed' runs N independent " "(build → encode → denoise) passes; slower but more " "robust if union-hull occasionally fails.") p.add_argument("--ensemble_mode", choices=["medoid", "consensus", "confidence"], default="medoid", help="How to combine the N ensemble runs. " "'medoid' (default) picks the run most similar to " "the others by Hungarian-matched distance. " "'confidence' picks the run with the highest mean " "validity logit. 'consensus' merges vertices and " "majority-votes edges (lost recall in earlier tests).") p.add_argument("--no_position_refine", action="store_true", help="Disable position refinement. By default, after the " "selector (medoid / confidence) picks a run, each " "vertex of that run is replaced with the per-axis " "aggregate (see --ensemble_refine_agg) of itself " "and its Hungarian-matched correspondents in the " "other N-1 runs (within --ensemble_merge_m). " "Topology unchanged; only vertex positions denoised.") p.add_argument("--ensemble_refine_agg", choices=["median", "mean", "wmean"], default="mean", help="Aggregator for position refinement. 'mean' " "(default) has lower variance under iid noise; " "'median' is robust to outlier correspondents; " "'wmean' is a softmax-of-logit weighted mean (each " "correspondent contributes proportional to " "exp(validity_logit), so high-confidence runs " "dominate). 'wmean' only kicks in when validity " "logits are non-trivial — falls back to plain mean " "when weights are uniform or missing.") p.add_argument("--ensemble_top_k", type=int, default=0, help="Soft-medoid: restrict the correspondent / voter set " "used by position refinement and edge voting to the " "K most central runs (lowest pairwise Hungarian " "distance). 0 (default) = use all N. Typical values " "for N=16: 5-8. The picked run (medoid) is always " "included.") p.add_argument("--no_edge_vote", action="store_true", default=True, help="Keep the picked medoid/confidence run's edges " "verbatim. This is the default.") p.add_argument("--edge_vote", action="store_false", dest="no_edge_vote", help="Enable edge majority voting: after selection each " "non-picked run's edges are relabelled into the picked " "run's vertex index space via the same Hungarian match " "used for position refinement, and edges with >= " "--ensemble_edge_vote_frac of votes survive.") p.add_argument("--ensemble_edge_vote_frac", type=float, default=0.5, help="Edge vote fraction (default 0.5 = strict majority). " "An edge (i,j) in picked-vertex space is kept iff at " "least ceil(M * frac) runs vote for it, where M is " "the active voter set (top_k if set, else N).") p.add_argument("--ensemble_vertex_vote_frac", type=float, default=0.0, help="Vertex consensus drop fraction (default 0.0 = off). " "When > 0, a picked-run vertex survives iff at least " "ceil(M * frac) voter runs have a Hungarian " "correspondent within tau. The picked run itself is " "always one voter (self-vote). Tightens precision by " "dropping spurious one-run-only detections — " "complements --ensemble_edge_vote. Used by both " "medoid/confidence modes (picked-then-refine) and " "consensus mode (sets the cluster-support threshold).") p.add_argument("--ensemble_stage2_last_k", type=int, default=1, help="Average the last K stage-2 denoiser steps' " "predictions (positions + validity / edge logits) " "before thresholding. 1 (default) = use only the " "final step (current behaviour). 3-5 mildly smooths " "ODE jitter near the trajectory endpoint at zero " "extra compute.") p.add_argument("--ensemble_selector", choices=["legacy", "ranker", "fixed"], default="legacy", help="Final ensemble chooser. 'legacy' (default) keeps the " "medoid/confidence/consensus path with no reranker. " "'ranker' builds " "raw/refined/voted/additive-edge candidates and picks " "with --ensemble_ranker_path if present, otherwise a " "small built-in heuristic. 'fixed' picks a named " "candidate family.") p.add_argument("--ensemble_fixed_candidate", default="confidence_mean_add", help="Candidate name used by --ensemble_selector fixed. " "Default confidence_mean_add, falling back to " "medoid_mean_add when the confidence candidate is " "identical to the medoid and therefore absent.") p.add_argument("--ensemble_topk_fuse", type=int, default=3, help="Top-K ranker-scored candidates to fuse (mean-refine " "positions + edge-vote union) when " "--ensemble_selector=ranker. 1 disables fusion and " "keeps the argmax candidate.") p.add_argument("--ensemble_ranker_path", default="ensemble_ranker_v5.npz", help="NPZ produced by --train_ensemble_ranker_jsonl. Used " "only when --ensemble_selector=ranker. If missing, " "ranker mode falls back to a heuristic scorer.") p.add_argument("--ensemble_candidate_dump", default=None, help="In --sanity mode, write ranker candidate features " "and GT HSS labels as JSONL for fast selector training.") p.add_argument("--train_ensemble_ranker_jsonl", default=None, help="Train the tiny NumPy ridge ranker from a candidate " "JSONL dump and exit. No model checkpoint/GPU needed.") p.add_argument("--ensemble_ranker_out", default="ensemble_ranker.npz", help="Output path for --train_ensemble_ranker_jsonl.") p.add_argument("--ensemble_ranker_l2", type=float, default=1e-2, help="Ridge L2 for --train_ensemble_ranker_jsonl.") # ---- local sanity check (no submission) ------------------------------ p.add_argument("--sanity", action="store_true", help="Local pre-submission sanity check: load default ckpt, " "stream a few val samples from the public trainval " "dataset, score against GT, print HSS, and exit. " "Writes no submission.json.") p.add_argument("--sanity_n", type=int, default=0, help="Number of val samples to score in --sanity mode. " "0 (default) means score the entire split.") p.add_argument("--sanity_split", default="validation", help="Split to draw sanity samples from (default 'validation').") p.add_argument("--sanity_dataset", default="usm3d/hoho22k_2026_trainval", help="HF dataset id with GT used for --sanity scoring.") p.add_argument("--sanity_cache_dir", default="cache", help="Local HF cache dir for sanity dataset streaming.") p.add_argument("--sanity_ensemble_diagnostics", action="store_true", help="In --sanity mode, score ensemble alternatives: " "member0, medoid pick, confidence pick, final fused, " "and oracle-best member. Diagnostic only; normal " "predictions are unchanged.") return p.parse_args(argv) def main(argv=None): args = parse_args(argv) if args.train_ensemble_ranker_jsonl: train_ensemble_ranker_jsonl( args.train_ensemble_ranker_jsonl, args.ensemble_ranker_out, l2=args.ensemble_ranker_l2, ) return # In --sanity mode we don't need /tmp/data, params.json, or the # anonymised test shards. Just resolve a checkpoint and score against # the public trainval GT. if args.sanity: params: Dict[str, Any] = {} if Path(args.params).exists(): with open(args.params) as f: params = json.load(f) ckpt_path = resolve_ckpt_path(args.ckpt, params) if args.device == "cuda" and not torch.cuda.is_available(): if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): device = torch.device("mps") else: device = torch.device("cpu") else: device = torch.device(args.device) if device.type == "cpu": print("[warn] CUDA/MPS unavailable — running on CPU; this will be slow.", flush=True) model, ckpt_args = load_model(ckpt_path, device) n_pts = (args.n_pts if args.n_pts is not None else ckpt_args.get("n_pts", 8192)) use_depth = ckpt_args.get("use_depth", True) n_sample_steps = (args.n_sample_steps if args.n_sample_steps is not None else ckpt_args.get("n_sample_steps", 50)) validity_thresh = (args.validity_thresh if args.validity_thresh is not None else ckpt_args.get("validity_thresh", 0.0)) stage2_cfg = _build_stage2_cfg(args, params, device) run_sanity_check( model=model, device=device, n_pts=n_pts, use_depth=use_depth, n_sample_steps=n_sample_steps, validity_thresh=validity_thresh, n_samples=args.sanity_n, split=args.sanity_split, dataset_id=args.sanity_dataset, cache_dir=args.sanity_cache_dir, stage2_cfg=stage2_cfg, ensemble_diagnostics=( args.sanity_ensemble_diagnostics or args.ensemble_candidate_dump is not None ), candidate_dump_path=args.ensemble_candidate_dump, ) return print("------------ Loading dataset ------------", flush=True) with open(args.params) as f: params = json.load(f) print(params, flush=True) print("pwd:", flush=True); os.system("pwd") os.system("ls -lahtr") print("/tmp/data/"); os.system("ls -lahtr /tmp/data/ 2>/dev/null") print("/tmp/data/data"); os.system("ls -lahtrR /tmp/data/data 2>/dev/null") dataset = load_test_dataset(params) print(dataset, flush=True) # ---- model: loaded ONCE ----------------------------------------------- ckpt_path = resolve_ckpt_path(args.ckpt, params) if args.device == "cuda" and not torch.cuda.is_available(): if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): device = torch.device("mps") else: device = torch.device("cpu") else: device = torch.device(args.device) if device.type == "cpu": print("[warn] CUDA/MPS unavailable — running on CPU; this will be slow.", flush=True) model, ckpt_args = load_model(ckpt_path, device) n_pts = (args.n_pts if args.n_pts is not None else ckpt_args.get("n_pts", 8192)) use_depth = ckpt_args.get("use_depth", True) n_sample_steps = (args.n_sample_steps if args.n_sample_steps is not None else ckpt_args.get("n_sample_steps", 50)) validity_thresh = (args.validity_thresh if args.validity_thresh is not None else ckpt_args.get("validity_thresh", 0.0)) print(f"[infer] n_pts={n_pts} use_depth={use_depth} " f"n_sample_steps={n_sample_steps} validity_thresh={validity_thresh} " f"n_prefetch={args.n_prefetch}", flush=True) stage2_cfg = _build_stage2_cfg(args, params, device) # ---- inference -------------------------------------------------------- print("------------ Running inference ------------", flush=True) solution: List[Dict[str, Any]] = [] for subset_name in dataset: print(f"[predict] subset {subset_name}", flush=True) split_ds = dataset[subset_name] if args.limit > 0: split_ds = split_ds.select(range(min(args.limit, len(split_ds)))) solution.extend( predict_split( split_ds, model=model, device=device, n_pts=n_pts, use_depth=use_depth, n_sample_steps=n_sample_steps, validity_thresh=validity_thresh, n_prefetch=args.n_prefetch, stage2_cfg=stage2_cfg, ) ) print("------------ Saving results ------------", flush=True) with open(args.output, "w") as f: json.dump(solution, f) print(f"[done] wrote {len(solution)} predictions → {args.output}", flush=True) if __name__ == "__main__": main()