data/preprocess.py

"""Scene feature point cloud builder from raw HoHo dataset samples.

V10 cache layout (additions over v7):
  - per-point RGB feature `scene_rgb` (uint8, 0..255):
      * COLMAP points use `point3D.color` from the COLMAP record
      * depth-unprojected points inherit the colour of their nearest COLMAP
        neighbour in raw-world space (no raw images are released)
      * camera tokens use a sentinel (0, 0, 0)
  - metric voxel downsampling across ALL candidate points (COLMAP + depth +
    camera) before the Tier-1/Tier-2 budget cap. At most one random point is
    kept per `voxel_size_m`-side voxel in raw world coordinates. Set
    `voxel_size_m <= 0` to disable.

Inherited from v7:
  - robust normalisation (median centre + P95 scale on COLMAP+camera union)
  - bbox_R kept identity (random yaw is applied as data augmentation in the loader)
  - separate per-point geometric and semantic confidence channels
  - soft top-2 gestalt labels (`scene_gestalt_ids`/`scene_gestalt_id2` + `scene_gestalt_w1`)
  - 3x3 majority gestalt/ADE projection on COLMAP points
  - fixed split: 4096 COLMAP/camera tokens + 4096 depth-unprojected tokens
  - camera tokens live only in the COLMAP half (type=2, Tier 0)
  - Tier 1 structural Gestalt classes are preferred before weaker Tier 2 classes
  - Tier 2 points require an ADE house/foreground label
  - hard skip on empty COLMAP / missing cam metadata / ambiguous image matches
"""
import io
from typing import Optional, Tuple

import numpy as np
from PIL import Image

from hoho2025.example_solutions import _cam_matrix_from_image, get_fitted_dense_depth
from color_mappings import ade20k_color_mapping, gestalt_color_mapping

# ------------------------------------------------------------------
# Gestalt / ADE20K palette — packed int32 LUTs built once at import time.
# ------------------------------------------------------------------
GESTALT_CLASSES = list(gestalt_color_mapping.keys())  # 28 classes, order = ID
ADE20K_CLASSES = list(ade20k_color_mapping.keys())


def _build_palette_lut(color_mapping, start_id: int = 0):
    """Return (packed_sorted, ids_sorted) arrays for searchsorted lookup."""
    colors = np.array([rgb for _, rgb in color_mapping.items()], dtype=np.int32)
    packed = (colors[:, 0] << 16) | (colors[:, 1] << 8) | colors[:, 2]
    ids = np.arange(start_id, start_id + len(packed), dtype=np.int64)
    order = np.argsort(packed)
    return packed[order], ids[order]


_GESTALT_PACKED_S, _GESTALT_IDS_S = _build_palette_lut(gestalt_color_mapping, start_id=0)
_ADE_PACKED_S, _ADE_IDS_S = _build_palette_lut(ade20k_color_mapping, start_id=1)


def _sample_palette(
    img_np: np.ndarray,
    vi: np.ndarray,
    ui: np.ndarray,
    packed_sorted: np.ndarray,
    ids_sorted: np.ndarray,
    default: int = 0,
) -> np.ndarray:
    """Look up palette class IDs for pixels at (vi, ui) without converting the whole image."""
    rgb = img_np[vi, ui].astype(np.int32)
    packed = (rgb[:, 0] << 16) | (rgb[:, 1] << 8) | rgb[:, 2]
    pos = np.searchsorted(packed_sorted, packed)
    in_range = pos < len(packed_sorted)
    hit = in_range & (packed_sorted[np.minimum(pos, len(packed_sorted) - 1)] == packed)
    return np.where(hit, ids_sorted[np.minimum(pos, len(ids_sorted) - 1)], default).astype(np.int64)


def gestalt_img_to_ids(gest_np: np.ndarray) -> np.ndarray:
    H, W = gest_np.shape[:2]
    vi, ui = np.mgrid[0:H, 0:W]
    return _sample_palette(
        gest_np,
        vi.ravel(),
        ui.ravel(),
        _GESTALT_PACKED_S,
        _GESTALT_IDS_S,
        default=0,
    ).reshape(H, W)


def ade20k_img_to_ids(ade_np: np.ndarray) -> np.ndarray:
    H, W = ade_np.shape[:2]
    vi, ui = np.mgrid[0:H, 0:W]
    return _sample_palette(
        ade_np,
        vi.ravel(),
        ui.ravel(),
        _ADE_PACKED_S,
        _ADE_IDS_S,
        default=0,
    ).reshape(H, W)


def _ids_from_names(names: tuple[str, ...], all_names: list[str], start_id: int = 0) -> frozenset:
    return frozenset(start_id + all_names.index(name) for name in names if name in all_names)


# Gestalt classes used in the fixed-budget scene cache.  The user-facing
# "flashing_end" and "transition line" names correspond to palette keys
# "flashing_end_point" and "transition_line" in this codebase.
GESTALT_TIER1_NAMES = (
    "apex",
    "eave_end_point",
    "flashing_end_point",
    "eave",
    "ridge",
    "rake",
    "valley",
)
GESTALT_TIER2_NAMES = (
    "ground_line",
    "flashing",
    "step_flashing",
    "hip",
    "fascia",
    "transition_line",
)

TIER1_GESTALT_IDS: frozenset = _ids_from_names(GESTALT_TIER1_NAMES, GESTALT_CLASSES)
TIER2_GESTALT_IDS: frozenset = _ids_from_names(GESTALT_TIER2_NAMES, GESTALT_CLASSES)
PRIORITY_GESTALT_IDS: frozenset = TIER1_GESTALT_IDS | TIER2_GESTALT_IDS

_TIER1_GESTALT_ARR = np.array(sorted(TIER1_GESTALT_IDS), dtype=np.int64)
_TIER2_GESTALT_ARR = np.array(sorted(TIER2_GESTALT_IDS), dtype=np.int64)
_PRIORITY_GESTALT_ARR = np.array(sorted(PRIORITY_GESTALT_IDS), dtype=np.int64)

ADE_HOUSE_FOREGROUND_NAMES = (
    "wall",
    "building;edifice",
    "windowpane;window",
    "door;double;door",
    "house",
    "column;pillar",
    "skyscraper",
    "stairs;steps",
)
ADE_HOUSE_FOREGROUND_IDS: frozenset = _ids_from_names(
    ADE_HOUSE_FOREGROUND_NAMES, ADE20K_CLASSES, start_id=1,
)
_ADE_HOUSE_FOREGROUND_ARR = np.array(sorted(ADE_HOUSE_FOREGROUND_IDS), dtype=np.int64)


def _ade_house_foreground_mask(ade_ids: np.ndarray) -> np.ndarray:
    return np.isin(ade_ids, _ADE_HOUSE_FOREGROUND_ARR)


# ------------------------------------------------------------------
# Image decode + COLMAP plumbing
# ------------------------------------------------------------------

def _image_to_rgb_array(img_like) -> np.ndarray:
    if isinstance(img_like, np.ndarray):
        arr = img_like
    elif isinstance(img_like, dict) and img_like.get("bytes") is not None:
        arr = np.array(Image.open(io.BytesIO(img_like["bytes"])).convert("RGB"))
    elif hasattr(img_like, "convert"):
        arr = np.array(img_like.convert("RGB"))
    else:
        arr = np.array(img_like)

    if arr.ndim == 2:
        arr = np.repeat(arr[..., None], 3, axis=-1)
    if arr.shape[-1] > 3:
        arr = arr[..., :3]
    return arr.astype(np.uint8, copy=False)


def _image_to_pil(img_like):
    """Decode image-like dataset values while preserving depth image modes."""
    if isinstance(img_like, Image.Image):
        return img_like
    if isinstance(img_like, dict) and img_like.get("bytes") is not None:
        img = Image.open(io.BytesIO(img_like["bytes"]))
        img.load()
        return img
    if isinstance(img_like, np.ndarray):
        return Image.fromarray(img_like)
    return img_like


def _camera_for_image(colmap_rec, col_img):
    try:
        return col_img.camera
    except AttributeError:
        return colmap_rec.cameras[col_img.camera_id]


def _observed_point_mask(col_img, point_ids: np.ndarray) -> np.ndarray:
    ids_attr = getattr(col_img, "point3D_ids", None)
    if ids_attr is not None:
        ids = ids_attr() if callable(ids_attr) else ids_attr
        ids = np.asarray(ids, dtype=np.int64)
        ids = ids[ids >= 0]
        if ids.size > 0:
            return np.isin(point_ids, ids)

    if hasattr(col_img, "has_point3D"):
        observed = np.fromiter(
            (bool(col_img.has_point3D(int(pid))) for pid in point_ids),
            dtype=bool,
            count=len(point_ids),
        )
        if observed.any():
            return observed

    return np.ones(len(point_ids), dtype=bool)


# ------------------------------------------------------------------
# COLMAP point + confidence extraction
# ------------------------------------------------------------------

def colmap_points_xyz(colmap_rec) -> np.ndarray:
    """Extract (N, 3) float32 world-space COLMAP points (xyz only)."""
    if not colmap_rec.points3D:
        return np.zeros((0, 3), dtype=np.float32)
    return np.array([p.xyz for p in colmap_rec.points3D.values()], dtype=np.float32)


def colmap_points_xyz_ids(colmap_rec):
    if not colmap_rec.points3D:
        return np.zeros(0, dtype=np.int64), np.zeros((0, 3), dtype=np.float32)
    ids, xyz = [], []
    for pid, p in colmap_rec.points3D.items():
        ids.append(int(pid))
        xyz.append(p.xyz)
    return np.asarray(ids, dtype=np.int64), np.asarray(xyz, dtype=np.float32)


def colmap_points_full(colmap_rec):
    """Return (ids, xyz, track_len, reproj_err, rgb) for all COLMAP 3D points.

    `rgb` is uint8 (N, 3) sourced from each point's COLMAP `color` field.
    """
    if not colmap_rec.points3D:
        return (
            np.zeros(0, dtype=np.int64),
            np.zeros((0, 3), dtype=np.float32),
            np.zeros(0, dtype=np.float32),
            np.zeros(0, dtype=np.float32),
            np.zeros((0, 3), dtype=np.uint8),
        )
    ids, xyz, tlen, err, rgb = [], [], [], [], []
    for pid, p in colmap_rec.points3D.items():
        ids.append(int(pid))
        xyz.append(p.xyz)
        try:
            tlen.append(int(p.track.length()))
        except Exception:
            try:
                tlen.append(int(len(p.track.elements)))
            except Exception:
                tlen.append(2)
        try:
            err.append(float(p.error))
        except Exception:
            err.append(1.0)
        try:
            rgb.append(np.asarray(p.color, dtype=np.uint8))
        except Exception:
            rgb.append(np.zeros(3, dtype=np.uint8))
    return (
        np.asarray(ids, dtype=np.int64),
        np.asarray(xyz, dtype=np.float32),
        np.asarray(tlen, dtype=np.float32),
        np.asarray(err, dtype=np.float32),
        np.asarray(rgb, dtype=np.uint8),
    )


def voxel_downsample_indices(
    xyz: np.ndarray,
    voxel_size: float,
    rng: np.random.Generator,
) -> np.ndarray:
    """Return indices into `xyz` keeping one random point per metric voxel.

    Voxel edges are aligned to a grid of side `voxel_size` in the same units
    as `xyz` (raw COLMAP world space). When `voxel_size <= 0` the function
    returns a permutation of all indices (no-op).
    """
    n = len(xyz)
    if n == 0 or voxel_size <= 0:
        return np.arange(n, dtype=np.int64)

    keys = np.floor(xyz / voxel_size).astype(np.int64)
    # lexsort by (z, y, x) so groups of identical voxel keys are contiguous
    order = np.lexsort(keys.T[::-1])
    keys_sorted = keys[order]

    diff = np.any(keys_sorted[1:] != keys_sorted[:-1], axis=1)
    starts = np.concatenate([[0], np.where(diff)[0] + 1])
    ends = np.concatenate([starts[1:], [n]])
    sizes = ends - starts
    offsets = rng.integers(0, sizes)  # one random pick per voxel
    return order[starts + offsets].astype(np.int64, copy=False)


def nearest_colmap_rgb(
    query_xyz: np.ndarray,
    colmap_xyz: np.ndarray,
    colmap_rgb: np.ndarray,
) -> np.ndarray:
    """For each query point, return the RGB of its nearest COLMAP neighbour.

    Falls back to all-zero RGB when COLMAP is empty. Uses scipy cKDTree when
    available; otherwise a chunked brute-force search.
    """
    n_q = len(query_xyz)
    if n_q == 0:
        return np.zeros((0, 3), dtype=np.uint8)
    if len(colmap_xyz) == 0:
        return np.zeros((n_q, 3), dtype=np.uint8)

    try:
        from scipy.spatial import cKDTree
        tree = cKDTree(colmap_xyz.astype(np.float32, copy=False))
        _, idx = tree.query(query_xyz.astype(np.float32, copy=False), k=1)
        idx = np.asarray(idx, dtype=np.int64).reshape(-1)
    except Exception:
        # Brute-force fallback in float32 chunks to bound memory.
        idx = np.zeros(n_q, dtype=np.int64)
        c = colmap_xyz.astype(np.float32, copy=False)
        chunk = 4096
        for s in range(0, n_q, chunk):
            q = query_xyz[s : s + chunk].astype(np.float32, copy=False)
            d2 = ((q[:, None, :] - c[None, :, :]) ** 2).sum(axis=-1)
            idx[s : s + chunk] = d2.argmin(axis=1)
    return colmap_rgb[idx]


def colmap_camera_centers(colmap_rec) -> np.ndarray:
    centers = []
    for img in colmap_rec.images.values():
        R, t = _cam_matrix_from_image(img)
        centers.append(R.T @ (-t))
    if not centers:
        return np.zeros((0, 3), dtype=np.float32)
    return np.array(centers, dtype=np.float32)


# ------------------------------------------------------------------
# Robust normalisation (no yaw alignment — that's a training augmentation now)
# ------------------------------------------------------------------

def _robust_norm_params(colmap_xyz: np.ndarray, cam_centers: np.ndarray):
    """Return (center, scale) for the scene. Rotation is identity by design."""
    pts = []
    if colmap_xyz.size:
        pts.append(colmap_xyz)
    if cam_centers.size:
        pts.append(cam_centers)
    pts = np.concatenate(pts, axis=0).astype(np.float64)

    center = np.median(pts, axis=0).astype(np.float32)
    d = np.linalg.norm(pts - center, axis=1)
    scale = float(max(np.percentile(d, 95.0), 1e-3))
    return center, scale


def _normalise(pts: np.ndarray, center: np.ndarray, scale: float) -> np.ndarray:
    return ((pts - center) / scale).astype(np.float32)


# ------------------------------------------------------------------
# Multi-view semantic projection (3x3 majority + soft top-2 + sem confidence)
# ------------------------------------------------------------------

def _project_to_pixel(pts: np.ndarray, col_img, cam):
    """Project (N, 3) world points to pixel space of a single camera."""
    cam_w = getattr(cam, "width", 0) or 0
    cam_h = getattr(cam, "height", 0) or 0
    if cam_w <= 0 or cam_h <= 0:
        return None

    R, t = _cam_matrix_from_image(col_img)
    p_cam = pts.astype(np.float64) @ R.T + t[None]
    z = p_cam[:, 2]
    in_front = z > 1e-8
    if not in_front.any():
        return None

    K = cam.calibration_matrix()
    u_proj = p_cam[:, 0] / np.maximum(z, 1e-8) * K[0, 0] + K[0, 2]
    v_proj = p_cam[:, 1] / np.maximum(z, 1e-8) * K[1, 1] + K[1, 2]
    return u_proj, v_proj, in_front, float(cam_w), float(cam_h)


def _gather_3x3_votes(
    img_np: np.ndarray,
    vi: np.ndarray,
    ui: np.ndarray,
    packed_sorted: np.ndarray,
    ids_sorted: np.ndarray,
    n_classes: int,
    default_id: int,
) -> np.ndarray:
    """Sample 3x3 neighbourhoods and return per-point class vote counts."""
    H, W = img_np.shape[:2]
    N = len(vi)
    counts = np.zeros((N, n_classes), dtype=np.uint16)
    for dv in (-1, 0, 1):
        for du in (-1, 0, 1):
            v2 = vi + dv
            u2 = ui + du
            ok = (v2 >= 0) & (v2 < H) & (u2 >= 0) & (u2 < W)
            if not ok.any():
                continue
            idx = np.where(ok)[0]
            cls = _sample_palette(
                img_np,
                v2[ok],
                u2[ok],
                packed_sorted,
                ids_sorted,
                default=default_id,
            )
            np.add.at(counts, (idx, cls), 1)
    return counts


def _resolve_colmap_image(colmap_by_name: dict, img_id: str):
    """Return the COLMAP image whose name matches img_id."""
    exact = colmap_by_name.get(img_id)
    if exact is not None:
        return exact
    matches = [v for k, v in colmap_by_name.items() if img_id in k]
    if len(matches) == 1:
        return matches[0]
    return None


def project_semantics_to_colmap_points(
    sample: dict,
    colmap_rec,
    point_ids: np.ndarray,
    point_xyz: np.ndarray,
) -> dict:
    """Vote projected Gestalt/ADE labels onto selected COLMAP 3D points (3x3)."""
    n_points = len(point_xyz)
    n_g = len(GESTALT_CLASSES)
    n_a = len(ADE20K_CLASSES) + 1
    gestalt_counts = np.zeros((n_points, n_g), dtype=np.uint32)
    ade_counts = np.zeros((n_points, n_a), dtype=np.uint32)

    image_ids = sample.get("image_ids", [])
    gestalt_imgs = sample.get("gestalt", [])
    ade_imgs = sample.get("ade", [])

    colmap_by_name = {col_img.name: col_img for col_img in colmap_rec.images.values()}

    for i, img_id in enumerate(image_ids):
        col_img = _resolve_colmap_image(colmap_by_name, img_id)
        if col_img is None:
            continue

        cam = _camera_for_image(colmap_rec, col_img)

        observed_mask = _observed_point_mask(col_img, point_ids)
        obs_idx = np.nonzero(observed_mask)[0]
        if obs_idx.size == 0:
            continue

        proj = _project_to_pixel(point_xyz[obs_idx], col_img, cam)
        if proj is None:
            continue
        u_proj, v_proj, in_front, cam_w, cam_h = proj

        def _valid_pixels(H: int, W: int):
            u = np.rint(u_proj * (W / cam_w)).astype(np.int64)
            v = np.rint(v_proj * (H / cam_h)).astype(np.int64)
            ok = in_front & (u >= 0) & (u < W) & (v >= 0) & (v < H)
            return u, v, ok

        if i < len(gestalt_imgs) and gestalt_imgs[i] is not None:
            gest_np = _image_to_rgb_array(gestalt_imgs[i])
            ui, vi, ok = _valid_pixels(*gest_np.shape[:2])
            if ok.any():
                votes = _gather_3x3_votes(
                    gest_np,
                    vi[ok],
                    ui[ok],
                    _GESTALT_PACKED_S,
                    _GESTALT_IDS_S,
                    n_classes=n_g,
                    default_id=0,
                )
                rows = obs_idx[ok]
                gestalt_counts[rows] += votes.astype(np.uint32)

        if i < len(ade_imgs) and ade_imgs[i] is not None:
            ade_np = _image_to_rgb_array(ade_imgs[i])
            ui, vi, ok = _valid_pixels(*ade_np.shape[:2])
            if ok.any():
                votes = _gather_3x3_votes(
                    ade_np,
                    vi[ok],
                    ui[ok],
                    _ADE_PACKED_S,
                    _ADE_IDS_S,
                    n_classes=n_a,
                    default_id=0,
                )
                rows = obs_idx[ok]
                ade_counts[rows] += votes.astype(np.uint32)

    gest_total = gestalt_counts.sum(axis=1)
    has_vote = gest_total > 0

    top1 = np.argmax(gestalt_counts, axis=1)
    counts2 = gestalt_counts.copy()
    counts2[np.arange(n_points), top1] = 0
    top2 = np.argmax(counts2, axis=1)
    has_second = counts2[np.arange(n_points), top2] > 0

    gest_id1 = np.where(has_vote, top1, -1).astype(np.int64)
    gest_id2 = np.where(has_vote & has_second, top2, -1).astype(np.int64)

    top1_count = gestalt_counts[np.arange(n_points), top1].astype(np.float32)
    gest_w1 = np.where(
        has_vote,
        top1_count / np.maximum(gest_total.astype(np.float32), 1.0),
        1.0,
    ).astype(np.float32)

    ade_ids = ade_counts.argmax(axis=1).astype(np.int64)

    return {
        "gestalt_id1": gest_id1,
        "gestalt_id2": gest_id2,
        "gestalt_w1": gest_w1,
        "ade_ids": ade_ids,
        "gest_top1_count": top1_count,
    }


# ------------------------------------------------------------------
# Tier-targeted depth unprojection (single-pixel labels)
# ------------------------------------------------------------------

def fit_depth_label_cache(
    sample: dict,
    colmap_rec,
) -> list[dict]:
    """Per-image preprocessing cache shared by stage-1 and stage-2 builders.

    For each image with valid depth fitting, returns the fitted depth map plus
    the gestalt + ADE label images resized to the depth resolution and the
    K/R/t needed to unproject pixels into world space.

    Stage 1's ``unproject_depth_tiered`` and stage 2's hull-cropped depth
    extractor both consume this cache, avoiding a second pass through
    ``get_fitted_dense_depth`` (Metric3Dv2) and a second PIL decode per image.
    """
    K_mats = np.array(sample["K"])
    R_mats = np.array(sample["R"])
    t_vecs = np.array(sample["t"])
    if t_vecs.ndim == 3:
        t_vecs = t_vecs[:, :, 0]
    pose_only = list(sample.get("pose_only_in_colmap", [False] * len(sample["image_ids"])))

    out: list[dict] = []
    for i, (depth_img, gest_img, ade_img, img_id) in enumerate(
        zip(sample["depth"], sample["gestalt"], sample["ade"], sample["image_ids"])
    ):
        if pose_only[i] or depth_img is None:
            continue
        Ki = K_mats[i]
        if Ki.shape != (3, 3) or Ki[0, 0] == 0:
            continue
        Ri = R_mats[i]
        ti = t_vecs[i]
        if not (np.isfinite(Ki).all() and np.isfinite(Ri).all() and np.isfinite(ti).all()):
            continue
        fx, fy = float(Ki[0, 0]), float(Ki[1, 1])
        if abs(fx) < 1e-6 or abs(fy) < 1e-6:
            continue

        depth_pil = _image_to_pil(depth_img)
        ade_pil   = _image_to_pil(ade_img) if ade_img is not None else None
        gest_pil  = _image_to_pil(gest_img) if gest_img is not None else None
        try:
            depth_fitted, _, found, _, _ = get_fitted_dense_depth(
                depth_pil, colmap_rec, img_id, ade_pil, verbose=False,
            )
        except Exception:
            continue
        if not found:
            continue

        H, W = depth_fitted.shape
        gest_np = (_image_to_rgb_array(gest_pil.resize((W, H), Image.NEAREST))
                   if gest_pil is not None else None)
        ade_np  = (_image_to_rgb_array(ade_pil.resize((W, H), Image.NEAREST))
                   if ade_pil is not None else None)
        out.append({
            "i":             int(i),
            "img_id":        img_id,
            "depth_fitted":  depth_fitted,
            "gest_np":       gest_np,
            "ade_np":        ade_np,
            "Ki":            Ki.astype(np.float32, copy=False),
            "Ri":            Ri.astype(np.float32, copy=False),
            "ti":            ti.astype(np.float32, copy=False),
        })
    return out


def unproject_depth_tiered(
    sample: dict,
    colmap_rec,
    n_per_image: int = 4096,
    rng: Optional[np.random.Generator] = None,
    depth_cache: Optional[list[dict]] = None,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Unproject Tier 1/2 Gestalt pixels. Returns (xyz, gestalt_ids, ade_ids, conf).

    Uses a single-pixel sample (no 3x3 voting): the pixel's own Gestalt/ADE IDs
    become the point labels. Tier 2 pixels are accepted only where ADE maps to
    the house/foreground subset, mirroring COLMAP token selection.

    When ``depth_cache`` (built by :func:`fit_depth_label_cache`) is supplied,
    the per-image PIL decode + Metric3Dv2 fitting is skipped — entries are
    consumed directly. Otherwise we build the cache implicitly.
    """
    if rng is None:
        rng = np.random.default_rng()

    if depth_cache is None:
        depth_cache = fit_depth_label_cache(sample, colmap_rec)

    all_xyz, all_gids, all_ade, all_conf = [], [], [], []

    for entry in depth_cache:
        if entry.get("gest_np") is None or entry.get("ade_np") is None:
            # Stage 1's tier filter needs both gestalt and ADE labels per pixel.
            continue
        depth_fitted = entry["depth_fitted"]
        gest_np      = entry["gest_np"]
        ade_np       = entry["ade_np"]
        Ki           = entry["Ki"]
        Ri           = entry["Ri"]
        ti           = entry["ti"]

        H, W = depth_fitted.shape
        gid_map = gestalt_img_to_ids(gest_np)
        ade_id_map = ade20k_img_to_ids(ade_np)

        valid_depth = depth_fitted > 0.1
        tier1_mask = np.isin(gid_map, _TIER1_GESTALT_ARR) & valid_depth
        tier2_mask = (
            np.isin(gid_map, _TIER2_GESTALT_ARR)
            & valid_depth
            & _ade_house_foreground_mask(ade_id_map)
        )

        def _sample_flat(mask: np.ndarray, budget: int) -> np.ndarray:
            if budget <= 0:
                return np.zeros(0, dtype=np.int64)
            flat = np.flatnonzero(mask.ravel())
            if flat.size > budget:
                flat = rng.choice(flat, budget, replace=False)
            return flat.astype(np.int64, copy=False)

        tier1_flat = _sample_flat(tier1_mask, n_per_image)
        tier2_flat = _sample_flat(tier2_mask, n_per_image - len(tier1_flat))
        if tier1_flat.size == 0 and tier2_flat.size == 0:
            continue

        sel = np.concatenate([tier1_flat, tier2_flat], axis=0)
        ys = (sel // W).astype(np.int64)
        xs = (sel %  W).astype(np.int64)
        depths = depth_fitted[ys, xs]

        fx, fy = Ki[0, 0], Ki[1, 1]
        cx, cy = Ki[0, 2], Ki[1, 2]
        x_cam = ((xs - cx) / fx) * depths
        y_cam = ((ys - cy) / fy) * depths
        z_cam = depths
        pts_cam = np.stack([x_cam, y_cam, z_cam], axis=1)
        pts_world = (pts_cam - ti[None]) @ Ri

        gids = gid_map[ys, xs].astype(np.int64)
        ade_ids_arr = ade_id_map[ys, xs].astype(np.int64)

        conf = np.clip(2.0 / np.maximum(depths, 0.5), 0.0, 1.0).astype(np.float32)

        all_xyz.append(pts_world.astype(np.float32))
        all_gids.append(gids)
        all_ade.append(ade_ids_arr.astype(np.int64))
        all_conf.append(conf)

    if not all_xyz:
        return (
            np.zeros((0, 3), dtype=np.float32),
            np.zeros(0, dtype=np.int64),
            np.zeros(0, dtype=np.int64),
            np.zeros(0, dtype=np.float32),
        )

    return (
        np.concatenate(all_xyz, axis=0),
        np.concatenate(all_gids, axis=0),
        np.concatenate(all_ade, axis=0),
        np.concatenate(all_conf, axis=0),
    )


def unproject_depth_priority(
    sample: dict,
    colmap_rec,
    n_per_image: int = 4096,
    rng: Optional[np.random.Generator] = None,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Backward-compatible alias for the tiered depth sampler."""
    return unproject_depth_tiered(sample, colmap_rec, n_per_image=n_per_image, rng=rng)


# ------------------------------------------------------------------
# Main scene-input builder
# ------------------------------------------------------------------

def _colmap_geom_conf(track_len: np.ndarray, reproj_err: np.ndarray) -> np.ndarray:
    """Geometric confidence of a COLMAP point in [0, 1]."""
    tl = np.clip(track_len / 5.0, 0.0, 1.0)
    er = np.exp(-np.maximum(reproj_err, 0.0) / 2.0)
    return (tl * er).astype(np.float32)


def build_scene_input(
    sample: dict,
    n_pts: int = 8192,
    use_depth: bool = True,
    n_depth_per_image: int = 4096,
    voxel_size_m: float = 0.0,
    rng: Optional[np.random.Generator] = None,
    colmap_sem: Optional[dict] = None,
    depth_cache: Optional[list[dict]] = None,
    return_cache: bool = False,
) -> dict:
    """Build a fixed-size scene feature point cloud from one dataset sample.

    `n_pts` is split in half by provenance (default 8192 = 4096 + 4096):
    the first half is COLMAP/camera tokens, the second half is depth-unprojected
    tokens. If a scene cannot supply enough points for either half, that half is
    padded by resampling from the points already selected for the same half.

    `voxel_size_m` (metres in raw world space) thins the candidate pool before
    the budget cap: COLMAP, depth and camera candidates are jointly voxelised
    and at most one random point per voxel survives. Set to 0 to disable.

    Tier order:
        COLMAP half:
            T0  camera centres                              (type=2)
            T1  apex/eave_end/flashing_end/eave/ridge/rake/valley
            T2  ground_line/flashing/step_flashing/hip/fascia/transition_line
                only when ADE is house/foreground
        Depth half:
            T1  same structural Gestalt classes             (type=1)
            T2  same weak Gestalt classes with house/foreground ADE
    """
    if rng is None:
        rng = np.random.default_rng()

    colmap_rec = sample["colmap"]

    colmap_pids, colmap_xyz, colmap_tlen, colmap_err, colmap_rgb = colmap_points_full(colmap_rec)
    cam_centers = colmap_camera_centers(colmap_rec)

    if len(colmap_xyz) == 0:
        raise ValueError("scene has no COLMAP 3D points — cannot build scene input")

    bbox_center, bbox_scale = _robust_norm_params(colmap_xyz, cam_centers)
    bbox_R = np.eye(3, dtype=np.float32)  # identity by design (yaw is augmentation)

    def norm(pts: np.ndarray) -> np.ndarray:
        return _normalise(pts, bbox_center, bbox_scale)

    # ---- COLMAP semantic projection ----------------------------------------
    # Reuse a precomputed semantic projection when one is supplied — the
    # 3x3-voting pass is the dominant per-sample cost in stage-1 preprocess.
    if colmap_sem is None:
        sem = project_semantics_to_colmap_points(sample, colmap_rec, colmap_pids, colmap_xyz)
    else:
        sem = colmap_sem
    all_g1  = sem["gestalt_id1"]
    all_g2  = sem["gestalt_id2"]
    all_w1  = sem["gestalt_w1"]
    all_ade = sem["ade_ids"]
    top1_count = sem["gest_top1_count"]

    geom_conf_all = _colmap_geom_conf(colmap_tlen, colmap_err)
    sem_conf_all  = np.clip(top1_count / 5.0, 0.0, 1.0).astype(np.float32)

    colmap_budget = n_pts // 2
    depth_budget = n_pts - colmap_budget

    tier1_mask = np.isin(all_g1, _TIER1_GESTALT_ARR)
    tier2_mask = np.isin(all_g1, _TIER2_GESTALT_ARR) & _ade_house_foreground_mask(all_ade)

    # ---- Depth unprojection (once) + nearest-COLMAP RGB --------------------
    if use_depth:
        # Build the per-image depth cache up here (instead of letting
        # unproject_depth_tiered build it internally) so we can capture it
        # for `return_cache=True` — otherwise stage 2 would redo Metric3Dv2
        # depth fitting per image, defeating the shared-compute refactor.
        if depth_cache is None:
            depth_cache = fit_depth_label_cache(sample, colmap_rec)
        d_xyz, d_gids, d_ade, d_conf = unproject_depth_tiered(
            sample, colmap_rec, n_per_image=n_depth_per_image, rng=rng,
            depth_cache=depth_cache,
        )
    else:
        d_xyz = np.zeros((0, 3), dtype=np.float32)
        d_gids = np.zeros(0, dtype=np.int64)
        d_ade = np.zeros(0, dtype=np.int64)
        d_conf = np.zeros(0, dtype=np.float32)

    # Build the COLMAP KDTree once; only query for depth points that survive
    # tier + voxel + budget selection. Avoids ~10-20x wasted queries on the
    # unprojected depth pool (which is typically ~80k while the final depth
    # half is ~4k).
    _colmap_kd = None
    if len(d_xyz) > 0 and len(colmap_xyz) > 0:
        try:
            from scipy.spatial import cKDTree
            _colmap_kd = cKDTree(colmap_xyz.astype(np.float32, copy=False))
        except Exception:
            _colmap_kd = None

    def _depth_rgb_for(sel: np.ndarray) -> np.ndarray:
        if len(sel) == 0:
            return np.zeros((0, 3), dtype=np.uint8)
        if _colmap_kd is None or len(colmap_rgb) == 0:
            return np.zeros((len(sel), 3), dtype=np.uint8)
        _, idx = _colmap_kd.query(d_xyz[sel].astype(np.float32, copy=False), k=1)
        return colmap_rgb[np.asarray(idx, dtype=np.int64).reshape(-1)]

    cam_rgb = np.zeros((len(cam_centers), 3), dtype=np.uint8)

    # ---- Joint voxel downsampling -----------------------------------------
    # Voxelise COLMAP + depth + camera candidates together in raw world space.
    # The kept-mask is sliced back per provenance and intersected with tier
    # masks below, so each Tier-1/Tier-2 selection picks from voxel survivors.
    if voxel_size_m and voxel_size_m > 0.0 and (len(colmap_xyz) + len(d_xyz) + len(cam_centers)) > 0:
        union_xyz = np.concatenate(
            [colmap_xyz, d_xyz, cam_centers.astype(np.float32, copy=False)], axis=0
        )
        keep_global = voxel_downsample_indices(union_xyz, float(voxel_size_m), rng)
        kept_flags = np.zeros(len(union_xyz), dtype=bool)
        kept_flags[keep_global] = True
        n_c = len(colmap_xyz)
        n_d = len(d_xyz)
        colmap_voxel_keep = kept_flags[:n_c]
        depth_voxel_keep = kept_flags[n_c : n_c + n_d]
        cam_voxel_keep = kept_flags[n_c + n_d :]
    else:
        colmap_voxel_keep = np.ones(len(colmap_xyz), dtype=bool)
        depth_voxel_keep = np.ones(len(d_xyz), dtype=bool)
        cam_voxel_keep = np.ones(len(cam_centers), dtype=bool)

    def _pick(mask: np.ndarray, budget: int) -> np.ndarray:
        if budget <= 0:
            return np.zeros(0, dtype=np.int64)
        idx = np.flatnonzero(mask)
        if len(idx) > budget:
            idx = rng.choice(idx, budget, replace=False)
        return idx.astype(np.int64, copy=False)

    def _empty_arrays(n: int = 0) -> tuple:
        return (
            np.zeros((n, 3), dtype=np.float32),   # 0: xyz
            np.zeros(n, dtype=np.int64),           # 1: type_id
            np.full(n, -1, dtype=np.int64),        # 2: gestalt_id1
            np.full(n, -1, dtype=np.int64),        # 3: gestalt_id2
            np.ones(n, dtype=np.float32),          # 4: gestalt_w1
            np.zeros(n, dtype=np.int64),           # 5: ade_id
            np.zeros(n, dtype=np.float32),         # 6: geom_conf
            np.zeros(n, dtype=np.float32),         # 7: sem_conf
            np.zeros((n, 3), dtype=np.uint8),      # 8: rgb
        )

    _N_FIELDS = 9

    def _concat(parts: list[tuple]) -> tuple:
        if not parts:
            return _empty_arrays(0)
        return tuple(np.concatenate([p[i] for p in parts], axis=0) for i in range(_N_FIELDS))

    def _take(arrs: tuple, idx: np.ndarray) -> tuple:
        return tuple(a[idx] for a in arrs)

    def _pad_or_trim(arrs: tuple, budget: int, fallback: Optional[tuple] = None) -> tuple:
        if budget <= 0:
            return _empty_arrays(0)
        n = len(arrs[0])
        if n == 0 and fallback is not None and len(fallback[0]) > 0:
            n_fallback = len(fallback[0])
            idx = rng.choice(n_fallback, budget, replace=n_fallback < budget)
            return _take(fallback, idx)
        if n == 0:
            return _empty_arrays(budget)
        if n > budget:
            idx = rng.choice(n, budget, replace=False)
            return _take(arrs, idx)
        if n < budget:
            pad = rng.choice(n, budget - n, replace=True)
            return tuple(np.concatenate([a, a[pad]], axis=0) for a in arrs)
        return arrs

    def _camera_arrays(sel: np.ndarray) -> tuple:
        n = len(sel)
        return (
            norm(cam_centers[sel]),
            np.full(n, 2, dtype=np.int64),
            np.full(n, -1, dtype=np.int64),
            np.full(n, -1, dtype=np.int64),
            np.ones(n, dtype=np.float32),
            np.zeros(n, dtype=np.int64),
            np.ones(n, dtype=np.float32),
            np.ones(n, dtype=np.float32),
            cam_rgb[sel],
        )

    def _colmap_arrays(sel: np.ndarray) -> tuple:
        n = len(sel)
        if n == 0:
            return _empty_arrays(0)
        return (
            norm(colmap_xyz[sel]),
            np.zeros(n, dtype=np.int64),
            all_g1[sel],
            all_g2[sel],
            all_w1[sel],
            all_ade[sel],
            geom_conf_all[sel],
            sem_conf_all[sel],
            colmap_rgb[sel],
        )

    def _depth_arrays(sel: np.ndarray) -> tuple:
        n = len(sel)
        if n == 0:
            return _empty_arrays(0)
        return (
            norm(d_xyz[sel]),
            np.full(n, 1, dtype=np.int64),
            d_gids[sel],
            np.full(n, -1, dtype=np.int64),
            np.ones(n, dtype=np.float32),
            d_ade[sel],
            d_conf[sel],
            np.ones(n, dtype=np.float32),
            _depth_rgb_for(sel),
        )

    # ---- COLMAP half: T0 cameras, then Tier 1, then Tier 2, then random fill
    col_parts: list[tuple] = []
    remaining_col = colmap_budget
    picked_colmap = np.zeros(len(colmap_xyz), dtype=bool)
    cam_eligible = np.flatnonzero(cam_voxel_keep)
    if remaining_col > 0 and len(cam_eligible) > 0:
        n_cam = min(len(cam_eligible), remaining_col)
        sel = (
            rng.choice(cam_eligible, n_cam, replace=False)
            if len(cam_eligible) > n_cam else cam_eligible
        )
        col_parts.append(_camera_arrays(sel.astype(np.int64, copy=False)))
        remaining_col -= n_cam

    if remaining_col > 0:
        sel = _pick(tier1_mask & colmap_voxel_keep & ~picked_colmap, remaining_col)
        col_parts.append(_colmap_arrays(sel))
        picked_colmap[sel] = True
        remaining_col -= len(sel)

    if remaining_col > 0:
        sel = _pick(tier2_mask & colmap_voxel_keep & ~picked_colmap, remaining_col)
        col_parts.append(_colmap_arrays(sel))
        picked_colmap[sel] = True
        remaining_col -= len(sel)

    # Fill whatever is left from ordinary COLMAP points (still voxel-thinned).
    if remaining_col > 0:
        sel = _pick(colmap_voxel_keep & ~picked_colmap, remaining_col)
        col_parts.append(_colmap_arrays(sel))
        picked_colmap[sel] = True
        remaining_col -= len(sel)

    col_arrays = _pad_or_trim(_concat(col_parts), colmap_budget)

    # ---- Depth half: Tier 1, then Tier 2 ------------------------------------
    depth_parts: list[tuple] = []
    if depth_budget > 0 and use_depth and len(d_xyz) > 0:
        d_tier1 = np.isin(d_gids, _TIER1_GESTALT_ARR) & depth_voxel_keep
        d_tier2 = (
            np.isin(d_gids, _TIER2_GESTALT_ARR)
            & _ade_house_foreground_mask(d_ade)
            & depth_voxel_keep
        )

        remaining_depth = depth_budget
        sel = _pick(d_tier1, remaining_depth)
        depth_parts.append(_depth_arrays(sel))
        remaining_depth -= len(sel)

        if remaining_depth > 0:
            sel = _pick(d_tier2, remaining_depth)
            depth_parts.append(_depth_arrays(sel))

    depth_arrays = _pad_or_trim(_concat(depth_parts), depth_budget, fallback=col_arrays)

    all_xyz       = np.concatenate([col_arrays[0], depth_arrays[0]], axis=0)
    all_type      = np.concatenate([col_arrays[1], depth_arrays[1]], axis=0)
    all_g1        = np.concatenate([col_arrays[2], depth_arrays[2]], axis=0)
    all_g2        = np.concatenate([col_arrays[3], depth_arrays[3]], axis=0)
    all_w1        = np.concatenate([col_arrays[4], depth_arrays[4]], axis=0)
    all_ade       = np.concatenate([col_arrays[5], depth_arrays[5]], axis=0)
    all_geom_conf = np.concatenate([col_arrays[6], depth_arrays[6]], axis=0)
    all_sem_conf  = np.concatenate([col_arrays[7], depth_arrays[7]], axis=0)
    all_rgb       = np.concatenate([col_arrays[8], depth_arrays[8]], axis=0)

    out = {
        "scene_xyz": all_xyz,
        "scene_type_ids": all_type,
        "scene_gestalt_ids": all_g1,
        "scene_gestalt_id2": all_g2,
        "scene_gestalt_w1": all_w1,
        "scene_ade_ids": all_ade,
        "scene_geom_conf": all_geom_conf.astype(np.float32),
        "scene_sem_conf":  all_sem_conf.astype(np.float32),
        "scene_rgb": all_rgb.astype(np.uint8),
        "bbox_center": bbox_center.astype(np.float32),
        "bbox_scale": float(bbox_scale),
        "bbox_R": bbox_R,
    }
    if return_cache:
        out["_cache"] = {
            "colmap": {
                "pids":      colmap_pids,
                "xyz":       colmap_xyz,
                "rgb":       colmap_rgb,
                "track_len": colmap_tlen,
                "reproj_err": colmap_err,
                "sem":       sem,
            },
            # Either the cache passed in, or the one built above for use_depth.
            # None only when use_depth=False (then stage 2 has no depth to share).
            "depth": depth_cache,
        }
    return out


# ------------------------------------------------------------------
# GT vertex builder
# ------------------------------------------------------------------

# ------------------------------------------------------------------
# V11 chimney filter — drop small connected components from GT
# ------------------------------------------------------------------

def _polygon_xy_area(verts_xy: np.ndarray) -> float:
    """Shoelace area of the convex polygon formed by 1..4 XY points.

    Vertices are first sorted by polar angle around their centroid so the
    polygon is non-self-intersecting. For < 3 points the area is zero.
    """
    n = len(verts_xy)
    if n < 3:
        return 0.0
    xy = verts_xy.astype(np.float64, copy=False)
    c = xy.mean(axis=0)
    ang = np.arctan2(xy[:, 1] - c[1], xy[:, 0] - c[0])
    order = np.argsort(ang)
    p = xy[order]
    x, y = p[:, 0], p[:, 1]
    return 0.5 * float(abs(np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1))))


def filter_small_components(
    wf_vertices: np.ndarray,
    wf_edges: np.ndarray,
    max_verts_per_cc: int = 4,
    max_area_m2: float = 5.0,
) -> Tuple[np.ndarray, np.ndarray]:
    """Mark vertices/edges belonging to small connected components for removal.

    Operates on RAW (un-normalised, metric) coordinates. A connected component
    is removed iff it has at most ``max_verts_per_cc`` vertices AND its XY
    footprint area is strictly less than ``max_area_m2`` square metres.

    Args:
        wf_vertices: (N, 3) world-space GT vertices.
        wf_edges:    (M, 2) edges as vertex-index pairs.

    Returns:
        (keep_v, keep_e): boolean masks of length N and M respectively. True
        means "retain" (not a chimney).
    """
    n = int(len(wf_vertices))
    m = int(len(wf_edges))
    if n == 0:
        return np.ones(0, dtype=bool), np.ones(m, dtype=bool)

    parent = np.arange(n, dtype=np.int64)

    def find(i: int) -> int:
        while parent[i] != i:
            parent[i] = parent[parent[i]]
            i = parent[i]
        return i

    if m > 0:
        edges_int = np.asarray(wf_edges, dtype=np.int64)
        valid_e = (
            (edges_int[:, 0] >= 0)
            & (edges_int[:, 1] >= 0)
            & (edges_int[:, 0] < n)
            & (edges_int[:, 1] < n)
        )
        for a, b in edges_int[valid_e]:
            ra, rb = find(int(a)), find(int(b))
            if ra != rb:
                parent[ra] = rb
    else:
        edges_int = np.zeros((0, 2), dtype=np.int64)
        valid_e = np.zeros(0, dtype=bool)

    roots = np.array([find(i) for i in range(n)], dtype=np.int64)

    keep_v = np.ones(n, dtype=bool)
    pts = np.asarray(wf_vertices, dtype=np.float32)
    for root in np.unique(roots):
        idx = np.where(roots == root)[0]
        if len(idx) <= max_verts_per_cc:
            area = _polygon_xy_area(pts[idx, :2])
            if area < max_area_m2:
                keep_v[idx] = False

    if m == 0:
        return keep_v, np.ones(0, dtype=bool)
    a = edges_int[:, 0].clip(0, n - 1)
    b = edges_int[:, 1].clip(0, n - 1)
    keep_e = valid_e & keep_v[a] & keep_v[b]
    return keep_v, keep_e


def build_gt_verts(
    sample: dict,
    bbox_center: np.ndarray,
    bbox_scale: float,
    k_verts: int = 64,
) -> np.ndarray:
    """Build padded GT vertex tensor of shape (K, 4) in normalised coords.

    Column 3 is the validity flag: +1 for real vertex, -1 for padding.
    No yaw rotation is applied here — the dataloader applies a random rotation
    to both `scene_xyz` and `verts_gt[:, :3]` per epoch.
    """
    verts = np.array(sample["wf_vertices"], dtype=np.float32)
    n_real = min(len(verts), k_verts)

    assert n_real > 0, "scene has zero GT vertices"

    out = np.zeros((k_verts, 4), dtype=np.float32)
    verts_norm = (verts[:n_real] - bbox_center) / bbox_scale
    out[:n_real, :3] = verts_norm.astype(np.float32)
    out[:n_real, 3] = 1.0
    out[n_real:, 3] = -1.0
    return out