Deploy learned baseline + hybrid multi-view tracking ensemble

base.json

@@ -0,0 +1,39 @@
+{
+    "arch": "perceiver",
+    "segments": 64,
+    "hidden": 256,
+    "ff": 1024,
+    "num_heads": 4,
+    "kv_heads_cross": 2,
+    "kv_heads_self": 2,
+    "latent_tokens": 256,
+    "latent_layers": 7,
+    "decoder_layers": 3,
+    "cross_attn_interval": 4,
+    "encoder_layers": 4,
+    "behind_emb_dim": 8,
+    "dropout": 0.1,
+    "activation": "gelu",
+    "rms_norm": true,
+    "qk_norm": true,
+    "qk_norm_type": "l2",
+    "segment_param": "midpoint_dir_len",
+    "segment_conf": true,
+    "vote_features": true,
+
+    "adam_betas": "0.9,0.95",
+    "weight_decay": 0.01,
+    "warmup": 10000,
+    "varifold_weight": 0.0,
+    "sinkhorn_weight": 1.0,
+    "sinkhorn_eps": 0.1,
+    "sinkhorn_iters": 20,
+    "sinkhorn_dustbin": 0.3,
+    "conf_weight": 0.1,
+    "conf_mode": "sinkhorn",
+    "conf_head_wd": 0.1,
+
+    "aug_rotate": true,
+    "aug_flip": true,
+    "seed": 353
+}

best_dgcnn_params.json

@@ -0,0 +1,8 @@
+{
+    "threshold": 0.6,
+    "strong_threshold": 0.7,
+    "very_strong_threshold": 0.85,
+    "max_length": 6.0,
+    "max_per_vertex": 1,
+    "dilate_px": 6
+}

evaluate_sklearn.py

@@ -0,0 +1,42 @@
+import io
+import json
+import numpy as np
+import sys
+import os
+from datasets import load_dataset
+from hoho2025.metric_helper import hss
+import sklearn_submission
+
+print("Loading dataset...")
+dataset = load_dataset('usm3d/hoho22k_2026_trainval', split='train', streaming=True, trust_remote_code=True)
+
+samples = []
+for idx, s in enumerate(dataset):
+    if idx >= 10:
+        break
+    samples.append(s)
+
+scores = []
+for idx, sample in enumerate(samples):
+    print(f"Testing sample {idx}")
+    try:
+        pred_v, pred_e = sklearn_submission.predict_wireframe_sklearn(sample)
+    except Exception as e:
+        print(f"Error on sample {idx}: {e}")
+        pred_v, pred_e = np.zeros((2, 3)), [(0, 1)]
+    
+    gt_v = sample.get('wf_vertices')
+    gt_e = sample.get('wf_edges')
+    
+    if gt_v is None or gt_e is None:
+        print(f"Skipping sample {idx} due to missing ground truth.")
+        continue
+
+    res = hss(pred_v, pred_e, gt_v, gt_e)
+    scores.append(res.hss)
+    print(f"Sample {idx} HSS: {res.hss:.4f}")
+
+if scores:
+    print(f"Average HSS: {sum(scores) / len(scores):.4f}")
+else:
+    print("No valid scores.")

s23dr_2026_example/attention.py

@@ -0,0 +1,141 @@
+# custom_transformer.py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# =============================================================================
+# Core Efficient Multihead Attention using Scaled Dot Product Attention (SDPA)
+# =============================================================================
+
+class MultiHeadSDPA(nn.Module):
+    """
+    Multi-head cross-attention using torch.nn.functional.scaled_dot_product_attention
+    without causal masking. Suitable for set inputs and cross-attention.
+
+    If qk_norm=True, L2-normalizes Q and K per-head before the dot product,
+    then scales by a learned per-head temperature (log_scale). This caps logit
+    magnitude to [-1, +1] * exp(log_scale), preventing attention entropy
+    collapse at large head_dim.
+    """
+    def __init__(self, d_model: int, num_heads: int, kv_heads: int = None,
+                 qk_norm: bool = False, qk_norm_type: str = "l2"):
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.kv_heads = kv_heads or num_heads
+        assert self.num_heads % self.kv_heads == 0, "kv_heads must divide num_heads"
+
+        self.head_dim = d_model // num_heads
+        self.qk_norm = qk_norm
+        self.qk_norm_type = qk_norm_type
+
+        # Input projection layers
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, self.kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(d_model, self.kv_heads * self.head_dim, bias=False)
+
+        # Output projection
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        nn.init.zeros_(self.out_proj.weight)
+
+        if qk_norm:
+            import math
+            if qk_norm_type == "rms":
+                # Standard QK-norm (Qwen3/Gemma3 style): RMSNorm on Q and K,
+                # no learned temperature. SDPA's 1/sqrt(d) scaling is sufficient
+                # because RMSNorm preserves the expected logit variance.
+                pass  # no extra parameters needed
+            else:
+                # L2 + learned temperature (nGPT/ViT-22B style):
+                # L2 projects to unit sphere, needs learned scale to compensate.
+                self.log_scale = nn.Parameter(
+                    torch.full((num_heads,), math.log(math.sqrt(self.head_dim))))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Project
+        q = self.q_proj(query)
+        k = self.k_proj(key)
+        v = self.v_proj(key)
+
+        B, Tq, _ = q.shape
+        _, Tk, _ = k.shape
+
+        q = q.view(B, Tq, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, Tk, self.kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, Tk, self.kv_heads, self.head_dim).transpose(1, 2)
+
+        if self.kv_heads != self.num_heads:
+            repeat = self.num_heads // self.kv_heads
+            k = k.repeat_interleave(repeat, dim=1)
+            v = v.repeat_interleave(repeat, dim=1)
+
+        if self.qk_norm:
+            if self.qk_norm_type == "rms":
+                # RMSNorm (Qwen3/Gemma3 style): no learned temperature needed.
+                # After RMSNorm, logit variance matches standard SDPA naturally.
+                q = q * torch.rsqrt(q.square().mean(dim=-1, keepdim=True) + 1e-6)
+                k = k * torch.rsqrt(k.square().mean(dim=-1, keepdim=True) + 1e-6)
+                attn_mask = None
+                if key_padding_mask is not None:
+                    attn_mask = ~key_padding_mask[:, None, None, :].to(dtype=torch.bool)
+                attn_out = F.scaled_dot_product_attention(
+                    q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False,
+                )
+            else:
+                # L2 + learned temperature (nGPT/ViT-22B style)
+                q = F.normalize(q, dim=-1)
+                k = F.normalize(k, dim=-1)
+                scale = self.log_scale.exp().view(1, -1, 1, 1)
+                q = q * scale
+                attn_mask = None
+                if key_padding_mask is not None:
+                    attn_mask = ~key_padding_mask[:, None, None, :].to(dtype=torch.bool)
+                attn_out = F.scaled_dot_product_attention(
+                    q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False,
+                    scale=1.0,
+                )
+        else:
+            attn_mask = None
+            if key_padding_mask is not None:
+                attn_mask = ~key_padding_mask[:, None, None, :].to(dtype=torch.bool)
+            attn_out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False
+            )
+
+        attn_out = attn_out.transpose(1, 2).reshape(B, Tq, self.d_model)
+        return self.out_proj(attn_out)
+
+
+# =============================================================================
+# Transformer Feed-Forward Block
+# =============================================================================
+
+def _get_activation(name: str):
+    """Look up activation function by name. Supports 'relu_sq' for ReLU^2."""
+    if name == "relu_sq":
+        return lambda x: F.relu(x).square()
+    return getattr(F, name)
+
+
+class FeedForward(nn.Module):
+    """
+    Position-wise MLP block: linear -> activation -> linear.
+    Supports 'gelu', 'relu', 'relu_sq', etc.
+    """
+    def __init__(self, d_model: int, dim_ff: int, activation: str = "gelu"):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, dim_ff)
+        self.linear2 = nn.Linear(dim_ff, d_model)
+        nn.init.zeros_(self.linear2.weight)
+        nn.init.zeros_(self.linear2.bias)
+        self.activation = _get_activation(activation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.linear1(x)
+        return self.linear2(self.activation(x))

s23dr_2026_example/bad_samples.txt

@@ -0,0 +1,156 @@
+14b1872e960
+1807ef90db4
+180e6a67e87
+1ad5c6bd31f
+1c3f939ad93
+1ede4c0d52f
+214f17d9cc4
+22256d88df9
+24a92a8de6d
+24b4e984bad
+2565978cf53
+2a71f1a2072
+2d44c1fade6
+2ebed43823a
+33982551420
+3b480496f82
+412a2bdf7a4
+44343bbabbb
+4a0b3f04cbd
+4a7fa170826
+4b7dc027214
+4e0dc2c9b18
+5172a516c8b
+529e8f15cd2
+56fc6f6f163
+575963ce814
+578ec40a278
+5a0c07c575a
+5d521223c26
+6148b5c9461
+631eb6d7c03
+655a14f8a75
+66502d7ee6f
+6da76fc6687
+777eaaad0ca
+7a4e2909d68
+7c5c9baf483
+80806dfd75e
+81a4ead431d
+833152dd554
+85797868c0f
+86460ad8181
+86783a6bee4
+95193322d7a
+99a9d056200
+9b1d4eeaab9
+9ff759f2e4c
+acbd243da16
+b9b275710c0
+beceaa9bb7c
+c243d079286
+c5c7337d2cb
+cdf6f2d3b35
+cfe370f1c87
+d4a72aea80c
+d655f066cd3
+d79e8d9455c
+d7d6c5be76e
+dc30ae4b93b
+de9495f7ca3
+e1901819c72
+e1d88c1a6b1
+e5d3eb0a617
+ec11d3cdcf6
+ecb21fad0ad
+ee55d8c6493
+ee7e6d4dee1
+008052054aa
+03ecb7d3cf3
+0555a655534
+099cad230c6
+0d061ae23f0
+10741a421c0
+110d5e407b9
+128a7fb415a
+13177736b26
+1635d73bf7d
+18a760de9ea
+18d90d03e95
+209627a5c1a
+21e3cd4b7b8
+22f5499200d
+266eb64de68
+269235f770b
+2758490e558
+2a203cf5d35
+2a878ec47ab
+2cb43eb2201
+393298e282b
+395abe6aac7
+3d19c7a4ca3
+44e2b719b1e
+45039819fcc
+4cb4ff01619
+4e5eb5712fa
+4e988765a6d
+5077bf42714
+55ed69b2622
+5ae3b651a37
+5ca1edeed4c
+5daa76b1c7f
+5fdd11dfae5
+6078cf180c2
+6682b309e9c
+6c02d2038c0
+71c595506c8
+73c8f960c18
+74ccc8fd057
+7a34156a798
+7ac7af9f59c
+7f2ec0ea179
+823b837b36c
+82d7600f9a3
+848161a2900
+88cedf129eb
+8dec106b6a6
+8e335d08ca4
+8ecf7c58193
+8fa55008beb
+90e09de2301
+9197acc0b9d
+954c25e876c
+98517d5563d
+99e717a0148
+9a0c0635bd7
+9ad436b7b3d
+9be351cbf14
+9e2a2e51798
+a84a7ea9220
+aa8cb84d3eb
+b07977292da
+b3e33456f0b
+b7823de373e
+bac379382d9
+bd2d9bf67a3
+c14584a84cd
+c497170c970
+cd8e767612b
+d17917bb279
+d42b9d432a9
+d53d8857a85
+d6808cf3d98
+d6f509d1dd9
+d7abd08e643
+d83493bf974
+d87293651ee
+da9d4ac9e8e
+daa1702791a
+dcb12411c14
+de9ab9cdd5b
+df906c58a3c
+e3870649eb5
+ea90aed9b98
+ecaa81b9711
+efc1238665b
+c5a65219daf

s23dr_2026_example/cache_scenes.py

@@ -0,0 +1,282 @@
+#!/usr/bin/env python3
+"""Cache compact scenes from HoHo22k shards to training-ready .pt files.
+
+Streams samples from the public `usm3d/hoho22k_2026_trainval` dataset, runs
+`build_compact_scene` (see point_fusion.py), precomputes priority group_id
+and semantic class_id, and saves one .pt per scene.
+
+Stage 1 of the dataset pipeline. See make_sampled_cache.py for stage 2.
+
+Usage:
+  python -m s23dr_2026_example.cache_scenes --out-dir cache/full --split train
+  python -m s23dr_2026_example.cache_scenes --out-dir cache/full_val --split validation
+
+Cache format per .pt file:
+  xyz:             float32 [P, 3]   all points in world space
+  source:          uint8   [P]      0=colmap, 1=depth
+  group_id:        int8    [P]      priority tier 0-4, -1=excluded
+  class_id:        uint8   [P]      one-hot class index (0-12)
+  behind_gest_id:  int16   [P]      behind-gestalt id (-1 if none)
+  visible_src:     uint8   [P]      1=gestalt, 2=ade
+  visible_id:      int16   [P]      class id within space
+  n_views_voted:   uint8   [P]      number of views that voted
+  vote_frac:       float32 [P]      fraction of votes
+  center:          float32 [3]      smart normalization center
+  scale:           float32 scalar   smart normalization scale
+  gt_vertices:     float32 [V, 3]   ground truth wireframe vertices
+  gt_edges:        int32   [E, 2]   ground truth wireframe edge indices
+"""
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .point_fusion import (
+    FuserConfig, build_compact_scene,
+    GEST_ID_TO_NAME, ADE_ID_TO_NAME, NUM_GEST,
+)
+
+# ---------------------------------------------------------------------------
+# Semantic class encoding: 11 structural + 1 other_house + 1 non_house = 13
+# ---------------------------------------------------------------------------
+
+# Each structural gestalt class gets its own one-hot bit.
+STRUCTURAL_CLASSES = (
+    "apex", "eave_end_point", "flashing_end_point",     # point classes (tier 0)
+    "rake", "ridge", "eave", "hip", "valley",           # roof edges (tier 1)
+    "flashing", "step_flashing",
+    "roof",                                              # roof face (tier 2)
+)
+# Index 11 = other house part (door, window, siding, etc.)
+# Index 12 = non-house / ADE / unlabeled
+NUM_SEMANTIC_CLASSES = len(STRUCTURAL_CLASSES) + 2  # 13
+
+# Priority tiers (same as tokenizer.py)
+_GEST_NAME_TO_ID = {n: i for i, n in enumerate(GEST_ID_TO_NAME)}
+_POINT_IDS = {_GEST_NAME_TO_ID[n] for n in ("apex", "eave_end_point", "flashing_end_point") if n in _GEST_NAME_TO_ID}
+_EDGE_IDS = {_GEST_NAME_TO_ID[n] for n in ("rake", "ridge", "eave", "hip", "valley", "flashing", "step_flashing") if n in _GEST_NAME_TO_ID}
+_FACE_IDS = {_GEST_NAME_TO_ID[n] for n in ("roof",) if n in _GEST_NAME_TO_ID}
+_HOUSE_IDS = {_GEST_NAME_TO_ID[n] for n in (
+    "apex", "eave_end_point", "flashing_end_point",
+    "rake", "ridge", "eave", "hip", "valley", "flashing", "step_flashing",
+    "roof", "door", "garage", "window", "shutter", "fascia", "soffit",
+    "horizontal_siding", "vertical_siding", "brick", "concrete",
+    "other_wall", "trim", "post", "ground_line",
+) if n in _GEST_NAME_TO_ID}
+
+_ADE_NAME_TO_ID = {n.lower(): i for i, n in enumerate(ADE_ID_TO_NAME)}
+_ADE_HOUSE_IDS = {_ADE_NAME_TO_ID[n] for n in ("building;edifice", "house", "wall", "windowpane;window", "door;double;door") if n in _ADE_NAME_TO_ID}
+
+_UNCLS_ID = _GEST_NAME_TO_ID.get("unclassified", -1)
+
+# Map structural gestalt names to one-hot index
+_STRUCTURAL_ONEHOT = {}
+for idx, name in enumerate(STRUCTURAL_CLASSES):
+    gid = _GEST_NAME_TO_ID.get(name)
+    if gid is not None:
+        _STRUCTURAL_ONEHOT[gid] = idx
+
+
+def _compute_group_and_class(visible_src, visible_id, behind_id, source):
+    """Compute priority group_id and semantic class_id per point (vectorized).
+
+    Args:
+        visible_src: uint8 [P] -- 0=unlabeled, 1=gestalt, 2=ade
+        visible_id:  int16 [P] -- class id within gestalt or ade space
+        behind_id:   int16 [P] -- behind-gestalt id (-1 if none)
+        source:      uint8 [P] -- 0=colmap, 1=depth
+
+    Returns:
+        group_id:  int8  [P] -- priority tier 0-4, -1 for excluded (unclassified)
+        class_id:  uint8 [P] -- one-hot class index 0-12
+    """
+    P = len(visible_src)
+    vsrc = visible_src.astype(np.int32)
+    vid = visible_id.astype(np.int32)
+    bid = behind_id.astype(np.int32)
+
+    # Effective gestalt id: prefer visible gestalt, fall back to behind
+    gest_id = np.full(P, -1, dtype=np.int32)
+    has_vis_gest = (vsrc == 1) & (vid >= 0)
+    has_behind = (bid >= 0) & ~has_vis_gest
+    gest_id[has_vis_gest] = vid[has_vis_gest]
+    gest_id[has_behind] = bid[has_behind]
+
+    # Exclude unclassified points
+    if _UNCLS_ID >= 0:
+        is_uncls = ((vsrc == 1) & (vid == _UNCLS_ID)) | (bid == _UNCLS_ID)
+        gest_id[is_uncls] = -1  # force excluded
+
+    # Build lookup arrays for gestalt id -> group and gestalt id -> class
+    max_gid = NUM_GEST
+    gid_to_group = np.full(max_gid, 4, dtype=np.int8)  # default: tier 4
+    gid_to_class = np.full(max_gid, NUM_SEMANTIC_CLASSES - 1, dtype=np.uint8)  # default: non-house
+
+    for gid in _POINT_IDS:
+        gid_to_group[gid] = 0
+    for gid in _EDGE_IDS:
+        gid_to_group[gid] = 1
+    for gid in _FACE_IDS:
+        gid_to_group[gid] = 2
+    for gid in _HOUSE_IDS - _POINT_IDS - _EDGE_IDS - _FACE_IDS:
+        gid_to_group[gid] = 3
+    for gid, onehot_idx in _STRUCTURAL_ONEHOT.items():
+        gid_to_class[gid] = onehot_idx
+    for gid in _HOUSE_IDS - set(_STRUCTURAL_ONEHOT.keys()):
+        gid_to_class[gid] = len(STRUCTURAL_CLASSES)  # other_house
+
+    # Apply lookup for points with valid gestalt ids
+    has_gest = gest_id >= 0
+    group_id = np.full(P, 4, dtype=np.int8)  # default: tier 4
+    class_id = np.full(P, NUM_SEMANTIC_CLASSES - 1, dtype=np.uint8)  # default: non-house
+
+    group_id[has_gest] = gid_to_group[gest_id[has_gest]]
+    class_id[has_gest] = gid_to_class[gest_id[has_gest]]
+
+    # ADE house points (no gestalt) get tier 3 + class_id = other_house
+    ade_house_arr = np.array(sorted(_ADE_HOUSE_IDS), dtype=np.int32)
+    is_ade_house = ~has_gest & (vsrc == 2) & (vid >= 0) & np.isin(vid, ade_house_arr)
+    group_id[is_ade_house] = 3
+    class_id[is_ade_house] = len(STRUCTURAL_CLASSES)  # other_house (index 11)
+
+    # Mark excluded points (unclassified) as -1
+    if _UNCLS_ID >= 0:
+        group_id[is_uncls] = -1
+        class_id[is_uncls] = NUM_SEMANTIC_CLASSES - 1
+
+    return group_id, class_id
+
+
+def _compute_smart_center_scale(xyz, source, mad_k=2.5, percentile=95.0,
+                                 max_points=8000):
+    """Compute normalization center and scale from depth points with MAD filter."""
+    depth_mask = source == 1
+    ref = xyz[depth_mask] if depth_mask.any() else xyz
+    if ref.shape[0] == 0:
+        center = xyz.mean(axis=0)
+        scale = max(np.linalg.norm(xyz - center, axis=1).max(), 1e-6)
+        return center.astype(np.float32), np.float32(scale)
+
+    if ref.shape[0] > max_points:
+        idx = np.random.choice(ref.shape[0], max_points, replace=False)
+        ref = ref[idx]
+
+    center0 = np.median(ref, axis=0)
+    dist = np.linalg.norm(ref - center0, axis=1)
+    med = np.median(dist)
+    mad = max(np.median(np.abs(dist - med)), 1e-6)
+    inliers = dist <= (med + mad_k * mad)
+    if inliers.any():
+        ref = ref[inliers]
+
+    # Percentile bounding box
+    lo_f = (100.0 - percentile) * 0.5 / 100.0
+    sorted_v = np.sort(ref, axis=0)
+    n = sorted_v.shape[0]
+    lo_idx = max(0, min(n - 1, int(lo_f * (n - 1))))
+    hi_idx = max(0, min(n - 1, int((1.0 - lo_f) * (n - 1))))
+    low = sorted_v[lo_idx]
+    high = sorted_v[hi_idx]
+
+    center = 0.5 * (low + high)
+    scale = max(np.sqrt(((high - low) ** 2).sum()), 1e-6)
+    return center.astype(np.float32), np.float32(scale)
+
+
+# ---------------------------------------------------------------------------
+# Dataset pipeline stage 1: raw HF sample -> cached .pt
+# ---------------------------------------------------------------------------
+
+def _process_one(sample, cfg):
+    """Fuse a single HF sample into a cache dict. Returns (order_id, dict) or None."""
+    rng = np.random.RandomState()
+
+    n_edges = len(sample.get("wf_edges", []))
+    if n_edges == 0 or n_edges > 64:
+        return None
+
+    scene = build_compact_scene(sample, cfg, rng=rng)
+    if scene is None:
+        return None
+
+    gt_v = scene.get("gt_vertices")
+    gt_e = scene.get("gt_edges")
+    if gt_v is None or gt_e is None or len(gt_e) == 0:
+        return None
+
+    xyz = scene["xyz"]
+    source = scene["source"]
+    group_id, class_id = _compute_group_and_class(
+        scene["visible_src"], scene["visible_id"], scene["behind_gest_id"], source)
+    center, scale = _compute_smart_center_scale(xyz, source)
+
+    gt_edge_classes = np.asarray(sample["wf_classifications"], dtype=np.int64)
+    return sample["order_id"], {
+        "xyz": xyz.astype(np.float32),
+        "source": source.astype(np.uint8),
+        "group_id": group_id,
+        "class_id": class_id,
+        "behind_gest_id": scene["behind_gest_id"].astype(np.int16),
+        "visible_src": scene["visible_src"].astype(np.uint8),
+        "visible_id": scene["visible_id"].astype(np.int16),
+        "n_views_voted": scene["n_views_voted"],
+        "vote_frac": scene["vote_frac"],
+        "center": center,
+        "scale": scale,
+        "gt_vertices": gt_v.astype(np.float32),
+        "gt_edges": gt_e.astype(np.int32),
+        "gt_edge_classes": gt_edge_classes,
+    }
+
+
+def main():
+    p = argparse.ArgumentParser(description="Stage 1: HoHo22k -> cached .pt files")
+    p.add_argument("--out-dir", required=True, help="Output directory for .pt files")
+    p.add_argument("--split", default="train", choices=["train", "validation"])
+    p.add_argument("--limit", type=int, default=0, help="Stop after N samples (0 = all)")
+    p.add_argument("--depth-per-view", type=int, default=8000)
+    p.add_argument("--skip-existing", action="store_true")
+    args = p.parse_args()
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    existing = {p.stem for p in out_dir.glob("*.pt")} if args.skip_existing else set()
+
+    from datasets import load_dataset
+    print(f"Streaming usm3d/hoho22k_2026_trainval split={args.split}...")
+    ds = load_dataset("usm3d/hoho22k_2026_trainval",
+                      streaming=True, trust_remote_code=True, split=args.split)
+
+    cfg = FuserConfig(depth_points_per_view=args.depth_per_view)
+    saved, skipped = 0, 0
+    t0 = time.perf_counter()
+    for i, sample in enumerate(ds):
+        if args.limit > 0 and i >= args.limit:
+            break
+        oid = sample["order_id"]
+        if oid in existing:
+            skipped += 1
+            continue
+        result = _process_one(sample, cfg)
+        if result is None:
+            skipped += 1
+            continue
+        order_id, data = result
+        torch.save(data, out_dir / f"{order_id}.pt")
+        saved += 1
+        if saved % 100 == 0:
+            rate = saved / (time.perf_counter() - t0)
+            print(f"  saved {saved} (skipped {skipped}) [{rate:.1f}/s]")
+
+    elapsed = time.perf_counter() - t0
+    print(f"Done. Saved {saved}, skipped {skipped} in {elapsed:.0f}s.")
+
+
+if __name__ == "__main__":
+    main()
+
+

s23dr_2026_example/color_mappings.py

@@ -0,0 +1,183 @@
+gestalt_color_mapping = {
+    "unclassified": (215, 62, 138),
+    "apex": (235, 88, 48),
+    "eave_end_point": (248, 130, 228),
+    "flashing_end_point": (71, 11, 161),
+    "ridge": (214, 251, 248),
+    "rake": (13, 94, 47),
+    "eave": (54, 243, 63),
+    "post": (187, 123, 236),
+    "ground_line": (136, 206, 14),
+    "flashing": (162, 162, 32),
+    "step_flashing": (169, 255, 219),
+    "hip": (8, 89, 52),
+    "valley": (85, 27, 65),
+    "roof": (215, 232, 179),
+    "door": (110, 52, 23),
+    "garage": (50, 233, 171),
+    "window": (230, 249, 40),
+    "shutter": (122, 4, 233),
+    "fascia": (95, 230, 240),
+    "soffit": (2, 102, 197),
+    "horizontal_siding": (131, 88, 59),
+    "vertical_siding": (110, 187, 198),
+    "brick": (171, 252, 7),
+    "concrete": (32, 47, 246),
+    "other_wall": (112, 61, 240),
+    "trim": (151, 206, 58),
+    "unknown": (127, 127, 127),
+    "transition_line": (0,0,0),
+}
+
+ade20k_color_mapping = {
+    'wall': (120, 120, 120),
+    'building;edifice': (180, 120, 120),
+    'sky': (6, 230, 230),
+    'floor;flooring': (80, 50, 50),
+    'tree': (4, 200, 3),
+    'ceiling': (120, 120, 80),
+    'road;route': (140, 140, 140),
+    'bed': (204, 5, 255),
+    'windowpane;window': (230, 230, 230),
+    'grass': (4, 250, 7),
+    'cabinet': (224, 5, 255),
+    'sidewalk;pavement': (235, 255, 7),
+    'person;individual;someone;somebody;mortal;soul': (150, 5, 61),
+    'earth;ground': (120, 120, 70),
+    'door;double;door': (8, 255, 51),
+    'table': (255, 6, 82),
+    'mountain;mount': (143, 255, 140),
+    'plant;flora;plant;life': (204, 255, 4),
+    'curtain;drape;drapery;mantle;pall': (255, 51, 7),
+    'chair': (204, 70, 3),
+    'car;auto;automobile;machine;motorcar': (0, 102, 200),
+    'water': (61, 230, 250),
+    'painting;picture': (255, 6, 51),
+    'sofa;couch;lounge': (11, 102, 255),
+    'shelf': (255, 7, 71),
+    'house': (255, 9, 224),
+    'sea': (9, 7, 230),
+    'mirror': (220, 220, 220),
+    'rug;carpet;carpeting': (255, 9, 92),
+    'field': (112, 9, 255),
+    'armchair': (8, 255, 214),
+    'seat': (7, 255, 224),
+    'fence;fencing': (255, 184, 6),
+    'desk': (10, 255, 71),
+    'rock;stone': (255, 41, 10),
+    'wardrobe;closet;press': (7, 255, 255),
+    'lamp': (224, 255, 8),
+    'bathtub;bathing;tub;bath;tub': (102, 8, 255),
+    'railing;rail': (255, 61, 6),
+    'cushion': (255, 194, 7),
+    'base;pedestal;stand': (255, 122, 8),
+    'box': (0, 255, 20),
+    'column;pillar': (255, 8, 41),
+    'signboard;sign': (255, 5, 153),
+    'chest;of;drawers;chest;bureau;dresser': (6, 51, 255),
+    'counter': (235, 12, 255),
+    'sand': (160, 150, 20),
+    'sink': (0, 163, 255),
+    'skyscraper': (140, 140, 140),
+    'fireplace;hearth;open;fireplace': (250, 10, 15),
+    'refrigerator;icebox': (20, 255, 0),
+    'grandstand;covered;stand': (31, 255, 0),
+    'path': (255, 31, 0),
+    'stairs;steps': (255, 224, 0),
+    'runway': (153, 255, 0),
+    'case;display;case;showcase;vitrine': (0, 0, 255),
+    'pool;table;billiard;table;snooker;table': (255, 71, 0),
+    'pillow': (0, 235, 255),
+    'screen;door;screen': (0, 173, 255),
+    'stairway;staircase': (31, 0, 255),
+    'river': (11, 200, 200),
+    'bridge;span': (255 ,82, 0),
+    'bookcase': (0, 255, 245),
+    'blind;screen': (0, 61, 255),
+    'coffee;table;cocktail;table': (0, 255, 112),
+    'toilet;can;commode;crapper;pot;potty;stool;throne': (0, 255, 133),
+    'flower': (255, 0, 0),
+    'book': (255, 163, 0),
+    'hill': (255, 102, 0),
+    'bench': (194, 255, 0),
+    'countertop': (0, 143, 255),
+    'stove;kitchen;stove;range;kitchen;range;cooking;stove': (51, 255, 0),
+    'palm;palm;tree': (0, 82, 255),
+    'kitchen;island': (0, 255, 41),
+    'computer;computing;machine;computing;device;data;processor;electronic;computer;information;processing;system': (0, 255, 173),
+    'swivel;chair': (10, 0, 255),
+    'boat': (173, 255, 0),
+    'bar': (0, 255, 153),
+    'arcade;machine': (255, 92, 0),
+    'hovel;hut;hutch;shack;shanty': (255, 0, 255),
+    'bus;autobus;coach;charabanc;double-decker;jitney;motorbus;motorcoach;omnibus;passenger;vehicle': (255, 0, 245),
+    'towel': (255, 0, 102),
+    'light;light;source': (255, 173, 0),
+    'truck;motortruck': (255, 0, 20),
+    'tower': (255, 184, 184),
+    'chandelier;pendant;pendent': (0, 31, 255),
+    'awning;sunshade;sunblind': (0, 255, 61),
+    'streetlight;street;lamp': (0, 71, 255),
+    'booth;cubicle;stall;kiosk': (255, 0, 204),
+    'television;television;receiver;television;set;tv;tv;set;idiot;box;boob;tube;telly;goggle;box': (0, 255, 194),
+    'airplane;aeroplane;plane': (0, 255, 82),
+    'dirt;track': (0, 10, 255),
+    'apparel;wearing;apparel;dress;clothes': (0, 112, 255),
+    'pole': (51, 0, 255),
+    'land;ground;soil': (0, 194, 255),
+    'bannister;banister;balustrade;balusters;handrail': (0, 122, 255),
+    'escalator;moving;staircase;moving;stairway': (0, 255, 163),
+    'ottoman;pouf;pouffe;puff;hassock': (255, 153, 0),
+    'bottle': (0, 255, 10),
+    'buffet;counter;sideboard': (255, 112, 0),
+    'poster;posting;placard;notice;bill;card': (143, 255, 0),
+    'stage': (82, 0, 255),
+    'van': (163, 255, 0),
+    'ship': (255, 235, 0),
+    'fountain': (8, 184, 170),
+    'conveyer;belt;conveyor;belt;conveyer;conveyor;transporter': (133, 0, 255),
+    'canopy': (0, 255, 92),
+    'washer;automatic;washer;washing;machine': (184, 0, 255),
+    'plaything;toy': (255, 0, 31),
+    'swimming;pool;swimming;bath;natatorium': (0, 184, 255),
+    'stool': (0, 214, 255),
+    'barrel;cask': (255, 0, 112),
+    'basket;handbasket': (92, 255, 0),
+    'waterfall;falls': (0, 224, 255),
+    'tent;collapsible;shelter': (112, 224, 255),
+    'bag': (70, 184, 160),
+    'minibike;motorbike': (163, 0, 255),
+    'cradle': (153, 0, 255),
+    'oven': (71, 255, 0),
+    'ball': (255, 0, 163),
+    'food;solid;food': (255, 204, 0),
+    'step;stair': (255, 0, 143),
+    'tank;storage;tank': (0, 255, 235),
+    'trade;name;brand;name;brand;marque': (133, 255, 0),
+    'microwave;microwave;oven': (255, 0, 235),
+    'pot;flowerpot': (245, 0, 255),
+    'animal;animate;being;beast;brute;creature;fauna': (255, 0, 122),
+    'bicycle;bike;wheel;cycle': (255, 245, 0),
+    'lake': (10, 190, 212),
+    'dishwasher;dish;washer;dishwashing;machine': (214, 255, 0),
+    'screen;silver;screen;projection;screen': (0, 204, 255),
+    'blanket;cover': (20, 0, 255),
+    'sculpture': (255, 255, 0),
+    'hood;exhaust;hood': (0, 153, 255),
+    'sconce': (0, 41, 255),
+    'vase': (0, 255, 204),
+    'traffic;light;traffic;signal;stoplight': (41, 0, 255),
+    'tray': (41, 255, 0),
+    'ashcan;trash;can;garbage;can;wastebin;ash;bin;ash-bin;ashbin;dustbin;trash;barrel;trash;bin': (173, 0, 255),
+    'fan': (0, 245, 255),
+    'pier;wharf;wharfage;dock': (71, 0, 255),
+    'crt;screen': (122, 0, 255),
+    'plate': (0, 255, 184),
+    'monitor;monitoring;device': (0, 92, 255),
+    'bulletin;board;notice;board': (184, 255, 0),
+    'shower': (0, 133, 255),
+    'radiator': (255, 214, 0),
+    'glass;drinking;glass': (25, 194, 194),
+    'clock': (102, 255, 0),
+    'flag': (92, 0, 255),
+}

s23dr_2026_example/data.py

@@ -0,0 +1,227 @@
+"""Data loading for pre-sampled HF datasets.
+
+Expects pre-sampled npz blobs with xyz_norm (not full PCD).
+Supports both 2048-point and 4096-point datasets.
+Use make_sampled_cache.py to produce these from full point clouds.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .tokenizer import EdgeDepthSequenceConfig
+
+# Default token budget (for 2048-point datasets; 4096 uses 3072/1024)
+SEQ_LEN = 2048
+COLMAP_POINTS = 1536
+DEPTH_POINTS = 512
+
+
+# ---------------------------------------------------------------------------
+# Datasets
+# ---------------------------------------------------------------------------
+
+def _load_bad_sample_ids():
+    """Load the set of known-bad sample IDs (misaligned GT, extreme scale)."""
+    bad_file = Path(__file__).parent / "bad_samples.txt"
+    if not bad_file.exists():
+        return set()
+    return set(line.strip() for line in bad_file.read_text().splitlines() if line.strip())
+
+
+class HFCachedDataset(torch.utils.data.Dataset):
+    """Load pre-sampled HuggingFace dataset into memory."""
+
+    def __init__(self, hf_dataset, aug_rotate=False, aug_jitter=0.0,
+                 aug_drop=0.0, aug_flip=False):
+        import io as _io
+        bad_ids = _load_bad_sample_ids()
+        print(f"Pre-decoding {len(hf_dataset)} samples into memory...")
+        self.samples = []
+        self.order_ids = []
+        n_skipped = 0
+        for i, sample in enumerate(hf_dataset):
+            if sample["order_id"] in bad_ids:
+                n_skipped += 1
+                continue
+            d = dict(np.load(_io.BytesIO(sample["data"])))
+            if "xyz_norm" not in d:
+                raise ValueError(
+                    f"Sample {sample['order_id']} missing 'xyz_norm' -- this looks like "
+                    f"a full PCD dataset, not pre-sampled. Use make_sampled_cache.py first.")
+            self.samples.append(d)
+            self.order_ids.append(sample["order_id"])
+            if (i + 1) % 2000 == 0:
+                print(f"  {i+1}/{len(hf_dataset)}...")
+        print(f"  Done. {len(self.samples)} samples in memory"
+              f" ({n_skipped} bad samples filtered).")
+        self.aug_rotate = aug_rotate
+        self.aug_jitter = aug_jitter
+        self.aug_drop = aug_drop
+        self.aug_flip = aug_flip
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        out = _process_sample(self.samples[idx], self.aug_rotate,
+                              self.aug_jitter, self.aug_drop, self.aug_flip)
+        out["sample_id"] = self.order_ids[idx]
+        return out
+
+
+def _process_sample(d, aug_rotate, aug_jitter=0.0, aug_drop=0.0, aug_flip=False):
+    """Process a pre-sampled npz dict into training tensors.
+
+    Args:
+        aug_rotate: random yaw rotation
+        aug_jitter: std of Gaussian noise added to point positions (0=disabled)
+        aug_drop: fraction of points to randomly drop (0=disabled)
+        aug_flip: random mirror along X axis (50% chance)
+    """
+    xyz_norm = d["xyz_norm"].copy()
+    gt_seg = d["gt_segments"].copy()
+    mask = d["mask"].copy()
+
+    if aug_rotate:
+        theta = np.random.rand() * 2 * np.pi
+        cos_t, sin_t = np.cos(theta), np.sin(theta)
+        x, z = xyz_norm[:, 0].copy(), xyz_norm[:, 2].copy()
+        xyz_norm[:, 0] = x * cos_t - z * sin_t
+        xyz_norm[:, 2] = x * sin_t + z * cos_t
+        for ep in range(2):
+            sx, sz = gt_seg[:, ep, 0].copy(), gt_seg[:, ep, 2].copy()
+            gt_seg[:, ep, 0] = sx * cos_t - sz * sin_t
+            gt_seg[:, ep, 2] = sx * sin_t + sz * cos_t
+
+    if aug_flip and np.random.rand() < 0.5:
+        xyz_norm[:, 0] = -xyz_norm[:, 0]
+        gt_seg[:, :, 0] = -gt_seg[:, :, 0]
+
+    if aug_jitter > 0:
+        valid = mask.astype(bool)
+        xyz_norm[valid] += np.random.randn(valid.sum(), 3).astype(np.float32) * aug_jitter
+
+    if aug_drop > 0:
+        valid_idx = np.where(mask)[0]
+        n_drop = int(len(valid_idx) * aug_drop)
+        if n_drop > 0:
+            drop_idx = np.random.choice(valid_idx, n_drop, replace=False)
+            mask[drop_idx] = False
+
+    result = {
+        "xyz_norm": torch.as_tensor(xyz_norm, dtype=torch.float32),
+        "class_id": torch.as_tensor(d["class_id"], dtype=torch.long),
+        "source": torch.as_tensor(d["source"], dtype=torch.long),
+        "mask": torch.as_tensor(mask),
+        "gt_segments": torch.as_tensor(gt_seg, dtype=torch.float32),
+        "scale": torch.tensor(float(d["scale"]), dtype=torch.float32),
+        "center": torch.as_tensor(d["center"], dtype=torch.float32),
+        "gt_vertices": d["gt_vertices"],
+        "gt_edges": d["gt_edges"],
+        "visible_src": torch.as_tensor(d["visible_src"], dtype=torch.long),
+        "visible_id": torch.as_tensor(d["visible_id"], dtype=torch.long),
+    }
+    if "behind" in d:
+        result["behind"] = torch.as_tensor(
+            np.clip(np.asarray(d["behind"], dtype=np.int16), 0, None), dtype=torch.long)
+    if "n_views_voted" in d:
+        result["n_views_voted"] = torch.as_tensor(d["n_views_voted"], dtype=torch.float32)
+    if "vote_frac" in d:
+        result["vote_frac"] = torch.as_tensor(d["vote_frac"], dtype=torch.float32)
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Collation + DataLoader
+# ---------------------------------------------------------------------------
+
+def collate(batch):
+    """Stack samples into batched tensors."""
+    out = {
+        "xyz_norm": torch.stack([d["xyz_norm"] for d in batch]),
+        "class_id": torch.stack([d["class_id"] for d in batch]),
+        "source": torch.stack([d["source"] for d in batch]),
+        "mask": torch.stack([d["mask"] for d in batch]),
+        "gt_segments": [d["gt_segments"] for d in batch],
+        "scales": torch.stack([d["scale"] for d in batch]),
+        "meta": batch,
+    }
+    # Optional fields: check ALL samples, not just batch[0].
+    # If any sample has it, all must have it (no mixed data versions).
+    for field in ("behind", "n_views_voted", "vote_frac"):
+        if any(field in d for d in batch):
+            missing = [i for i, d in enumerate(batch) if field not in d]
+            if missing:
+                raise KeyError(
+                    f"Field '{field}' present in some batch samples but missing in "
+                    f"{len(missing)}/{len(batch)}. Mixed data versions in cache?")
+            out[field] = torch.stack([d[field] for d in batch])
+    return out
+
+
+def build_loader(cache_dir, batch_size, aug_rotate=False, aug_jitter=0.0,
+                 aug_drop=0.0, aug_flip=False):
+    """Create a DataLoader from HF dataset.
+
+    cache_dir should be 'hf://repo/name:split' format.
+    """
+    if not cache_dir.startswith("hf://"):
+        raise ValueError(
+            f"cache_dir must be 'hf://repo:split' format, got: {cache_dir}. "
+            f"Local .pt caches are no longer supported in the training path.")
+    parts = cache_dir[5:].split(":")
+    repo = parts[0]
+    split = parts[1] if len(parts) > 1 else "train"
+    from datasets import load_dataset
+    hf_ds = load_dataset(repo, split=split)
+    ds = HFCachedDataset(hf_ds, aug_rotate=aug_rotate, aug_jitter=aug_jitter,
+                          aug_drop=aug_drop, aug_flip=aug_flip)
+    loader = torch.utils.data.DataLoader(
+        ds, batch_size=batch_size, shuffle=True,
+        num_workers=0, collate_fn=collate,
+    )
+    print(f"Dataset: {len(ds)} scenes, batch_size={batch_size}")
+    return loader
+
+
+# ---------------------------------------------------------------------------
+# Token building (GPU)
+# ---------------------------------------------------------------------------
+
+def build_tokens(batch, model, device):
+    """Apply Fourier features + learned embeddings on GPU."""
+    xyz = batch["xyz_norm"].to(device)
+    cid = batch["class_id"].to(device)
+    src = batch["source"].to(device)
+    masks = batch["mask"].to(device)
+    gt = [g.to(device) for g in batch["gt_segments"]]
+    scales = batch["scales"]
+
+    B, T, _ = xyz.shape
+    tok = model.tokenizer
+    fourier = tok.pos_enc(xyz.reshape(-1, 3)).reshape(B, T, -1) \
+        if tok.pos_enc is not None else xyz.new_zeros(B, T, 0)
+    parts = [xyz, fourier, tok.label_emb(cid), tok.src_emb(src.clamp(0, 1))]
+    if tok.behind_emb_dim > 0:
+        if "behind" in batch:
+            beh = batch["behind"].to(device)
+        else:
+            # Data doesn't have behind -- use zeros (embed index 0).
+            # This is intentional for eval on old data; for training,
+            # fail fast by requiring the field (checked in _process_sample).
+            beh = xyz.new_zeros(B, T, dtype=torch.long)
+        parts.append(tok.behind_emb(beh))
+    if tok.use_vote_features:
+        if "n_views_voted" not in batch or "vote_frac" not in batch:
+            raise KeyError(
+                "Model expects vote features (--vote-features) but data is missing "
+                "'n_views_voted'/'vote_frac'. Use v2 dataset or regenerate cache.")
+        # Normalize to ~zero mean, unit variance (dataset stats: nv~2.7+/-1.0, vf~0.5+/-0.25)
+        nv = ((batch["n_views_voted"].to(device).float() - 2.7) / 1.0).unsqueeze(-1)
+        vf = ((batch["vote_frac"].to(device).float() - 0.5) / 0.25).unsqueeze(-1)
+        parts.extend([nv, vf])
+    tokens = torch.cat(parts, dim=-1)
+    return tokens, masks, gt, scales, batch["meta"]

s23dr_2026_example/losses.py

@@ -0,0 +1,215 @@
+"""Loss computation for wireframe prediction."""
+from __future__ import annotations
+
+import torch
+
+from .varifold import varifold_loss_batch
+from .sinkhorn import batched_sinkhorn_loss
+
+# Varifold config
+VARIANT = "simpson3"
+SIGMAS = [0.5, 1.0, 2.0]  # meters (divided by per-scene scale at runtime)
+ALPHAS = [0.2, 0.6, 0.2]
+LEN_POW = 1.0
+VARIFOLD_CROSS_ONLY = False  # Set to True to drop self-energy (avoids O(S^2) blowup)
+
+# Sinkhorn config (note: near-zero gradients at eps=0.05, effectively disabled)
+SINKHORN_EPS = 0.05
+SINKHORN_ITERS = 10
+
+# Sinkhorn dustbin cost: controls the OT "not matching" penalty.
+# Like tau, this is an OT behavior parameter, NOT a physical distance.
+# Must be comparable to typical matching costs in normalized space (~0.1).
+# Do NOT divide by scale.
+SINKHORN_DUSTBIN = 0.1
+
+MAX_GT = 64  # fixed pad size for compile-friendly shapes
+
+# Precomputed constants (created once on first call)
+_loss_constants = {}
+
+
+def _get_loss_constants(device, dtype):
+    key = (device, dtype)
+    if key not in _loss_constants:
+        _loss_constants[key] = {
+            "sigmas": torch.tensor(SIGMAS, device=device, dtype=dtype),
+            "alphas": torch.tensor(ALPHAS, device=device, dtype=dtype),
+        }
+    return _loss_constants[key]
+
+
+def pad_gt_fixed(gt_list, device, dtype):
+    """Pad GT segments to fixed MAX_GT for compile-friendly shapes."""
+    B = len(gt_list)
+    gt_pad = torch.zeros((B, MAX_GT, 2, 3), device=device, dtype=dtype)
+    gt_mask = torch.zeros((B, MAX_GT), device=device, dtype=torch.bool)
+    gt_lengths = torch.zeros(B, device=device, dtype=dtype)
+    for i, g in enumerate(gt_list):
+        n = g.shape[0]
+        if n > 0:
+            gt_pad[i, :n] = g
+            gt_mask[i, :n] = True
+            gt_lengths[i] = torch.linalg.norm(g[:, 1] - g[:, 0], dim=-1).sum()
+    return gt_pad, gt_mask, gt_lengths
+
+
+def _loss_inner(pred_segments, gt_pad, gt_mask, gt_lengths, scales,
+                sigmas, alphas, varifold_w):
+    """Pure tensor loss -- no Python control flow, no boolean indexing."""
+    has_gt = (gt_lengths > 0).float()
+
+    sigmas_eff = sigmas / scales[:, None]
+    loss_batch = varifold_loss_batch(
+        pred_segments, gt_pad, gt_mask=gt_mask,
+        variant=VARIANT, sigmas=sigmas_eff, alpha=alphas, len_pow=LEN_POW,
+        cross_only=VARIFOLD_CROSS_ONLY,
+    )
+    v = loss_batch / gt_lengths.clamp(min=1.0)
+    v = (v * has_gt).sum() / has_gt.sum().clamp(min=1.0)
+
+    total = varifold_w * v
+    return total, v
+
+
+# Will be replaced with compiled version on CUDA
+_loss_fn = _loss_inner
+
+
+def compute_loss(pred_segments, gt_list, scales, device,
+                 varifold_w, sinkhorn_w,
+                 endpoint_w=0.0,
+                 conf_logits=None, conf_weight=0.0, conf_mode="sinkhorn",
+                 sinkhorn_eps=None, sinkhorn_iters=None,
+                 sinkhorn_dustbin=None, conf_clamp_min=None):
+    """Combined loss with fixed-size GT padding.
+
+    conf_mode: "sinkhorn" = conf-weighted sinkhorn, "sinkhorn_detach" = detached conf.
+    """
+    if conf_logits is not None and conf_clamp_min is not None:
+        conf_logits = conf_logits.clamp(min=conf_clamp_min)
+    gt_pad, gt_mask, gt_lengths = pad_gt_fixed(gt_list, device, pred_segments.dtype)
+    c = _get_loss_constants(device, pred_segments.dtype)
+
+    total, v = _loss_fn(
+        pred_segments, gt_pad, gt_mask, gt_lengths, scales,
+        c["sigmas"], c["alphas"], varifold_w)
+
+    terms = {}
+    if varifold_w > 0:
+        terms["varifold"] = v.detach()
+
+    if sinkhorn_w > 0:
+        has_gt = (gt_lengths > 0).float()
+        if conf_logits is not None and conf_mode == "sinkhorn":
+            pred_mass = torch.sigmoid(conf_logits)
+        elif conf_logits is not None and conf_mode == "sinkhorn_detach":
+            pred_mass = torch.sigmoid(conf_logits.detach())
+        else:
+            pred_mass = None
+        eps = sinkhorn_eps if sinkhorn_eps is not None else SINKHORN_EPS
+        iters = sinkhorn_iters if sinkhorn_iters is not None else SINKHORN_ITERS
+        dustbin = sinkhorn_dustbin if sinkhorn_dustbin is not None else SINKHORN_DUSTBIN
+        S = pred_segments.shape[1]
+        sink_per = batched_sinkhorn_loss(
+            pred_segments, gt_pad, gt_mask,
+            eps, iters, dustbin,
+            pred_mass=pred_mass,
+        ) / (gt_lengths.clamp(min=1.0) * S)
+        s = (sink_per * has_gt).sum() / has_gt.sum().clamp(min=1.0)
+        total = total + sinkhorn_w * s
+        terms["sinkhorn"] = s.detach()
+
+    if conf_logits is not None and conf_weight > 0:
+        if conf_mode in ("sinkhorn", "sinkhorn_detach"):
+            conf_w = torch.sigmoid(conf_logits)
+            S = conf_logits.shape[1]
+            gt_counts = gt_mask.sum(dim=1).float()
+            conf_sum = conf_w.sum(dim=1)
+            reg = (((conf_sum - gt_counts) / S) ** 2).mean()
+            total = total + conf_weight * reg
+            terms["conf_reg"] = reg.detach()
+        else:
+            raise ValueError(f"Unknown conf_mode: {conf_mode}")
+
+    if endpoint_w > 0:
+        has_gt = (gt_lengths > 0).float()
+        eps_ep = sinkhorn_eps if sinkhorn_eps is not None else SINKHORN_EPS
+        iters_ep = sinkhorn_iters if sinkhorn_iters is not None else SINKHORN_ITERS
+        dustbin_ep = sinkhorn_dustbin if sinkhorn_dustbin is not None else SINKHORN_DUSTBIN
+        B, S = pred_segments.shape[:2]
+        M = gt_pad.shape[1]
+
+        # Compute hard assignment via sinkhorn (detached -- matching is not trained)
+        with torch.no_grad():
+            pred_mass_ep = torch.sigmoid(conf_logits) if conf_logits is not None else None
+            sink_loss_for_assign = batched_sinkhorn_loss(
+                pred_segments, gt_pad, gt_mask, eps_ep, iters_ep, dustbin_ep,
+                pred_mass=pred_mass_ep)
+            p0, p1 = pred_segments[:, :, 0], pred_segments[:, :, 1]
+            g0, g1 = gt_pad[:, :, 0], gt_pad[:, :, 1]
+            mid_p, half_p = 0.5 * (p0 + p1), 0.5 * (p1 - p0)
+            mid_g, half_g = 0.5 * (g0 + g1), 0.5 * (g1 - g0)
+            d_mid = torch.linalg.norm(mid_p.unsqueeze(2) - mid_g.unsqueeze(1), dim=-1)
+            len_p = torch.linalg.norm(half_p, dim=-1, keepdim=True).clamp(min=1e-6)
+            len_g = torch.linalg.norm(half_g, dim=-1, keepdim=True).clamp(min=1e-6)
+            dir_p, dir_g = half_p / len_p, half_g / len_g
+            cos_a = (dir_p.unsqueeze(2) * dir_g.unsqueeze(1)).sum(dim=-1)
+            d_dir = 1.0 - cos_a.abs()
+            d_len = (len_p.unsqueeze(2) - len_g.unsqueeze(1)).squeeze(-1).abs()
+            cost = d_mid + d_dir + d_len
+            dc = torch.as_tensor(dustbin_ep, device=cost.device, dtype=cost.dtype)
+            cost = torch.where(gt_mask.unsqueeze(1), cost, dc * 10.0)
+            cost_pad = dc.expand(B, S + 1, M + 1).clone()
+            cost_pad[:, :S, :M] = cost
+            cost_pad[:, -1, -1] = 0.0
+            gt_counts = gt_mask.sum(dim=1).float()
+            if pred_mass_ep is not None:
+                pm = pred_mass_ep.clamp(min=0.0)
+                a = torch.cat([pm, (gt_counts - pm.sum(1)).clamp(min=0).unsqueeze(1)], dim=1)
+                b_val = torch.zeros(B, M + 1, device=cost.device, dtype=cost.dtype)
+                b_val[:, :M] = gt_mask.float()
+                b_val[:, -1] = (pm.sum(1) - gt_counts).clamp(min=0)
+            else:
+                n = float(S)
+                denom = n + gt_counts
+                a = (1.0 / denom).unsqueeze(1).expand(B, S + 1).clone()
+                a[:, -1] = gt_counts / denom
+                b_val = (1.0 / denom).unsqueeze(1).expand(B, M + 1).clone()
+                b_val[:, -1] = n / denom
+                b_val[:, :M] = b_val[:, :M] * gt_mask.float()
+            log_a = torch.log(a + 1e-9)
+            log_b = torch.log(b_val + 1e-9)
+            log_k = -cost_pad / eps_ep
+            log_u = torch.zeros_like(a)
+            log_v = torch.zeros_like(b_val)
+            for _ in range(iters_ep):
+                log_u = log_a - torch.logsumexp(log_k + log_v.unsqueeze(1), dim=2)
+                log_v = log_b - torch.logsumexp(log_k + log_u.unsqueeze(2), dim=1)
+            transport = torch.exp(log_u.unsqueeze(2) + log_v.unsqueeze(1) + log_k)
+            assignment = transport[:, :S, :M+1].argmax(dim=2)
+            assignment[assignment >= M] = -1
+
+        # Everything below is WITH gradients (assignment is detached but pred_segments is live)
+        matched = (assignment >= 0)  # [B, S]
+        n_matched = matched.float().sum().clamp(min=1.0)
+        assign_safe = assignment.clamp(min=0)
+        gt_matched = gt_pad[
+            torch.arange(B, device=device)[:, None].expand(B, S),
+            assign_safe]  # [B, S, 2, 3]
+
+        # Symmetric endpoint distance
+        ref_ep1 = pred_segments[:, :, 0]
+        ref_ep2 = pred_segments[:, :, 1]
+        gt_ep1 = gt_matched[:, :, 0]
+        gt_ep2 = gt_matched[:, :, 1]
+        dist_fwd = (ref_ep1 - gt_ep1).norm(dim=-1) + (ref_ep2 - gt_ep2).norm(dim=-1)
+        dist_rev = (ref_ep1 - gt_ep2).norm(dim=-1) + (ref_ep2 - gt_ep1).norm(dim=-1)
+        ep_dist = torch.min(dist_fwd, dist_rev)
+
+        # Normalize by GT total length * S (same scale as sinkhorn)
+        ep_loss = (ep_dist * matched.float()).sum() / n_matched
+        total = total + endpoint_w * ep_loss
+        terms["endpoint"] = ep_loss.detach()
+
+    return total, terms

s23dr_2026_example/make_sampled_cache.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+"""Stage 2: priority-sample cached .pt scenes into fixed-size .npz files.
+
+Reads the per-scene .pt files produced by cache_scenes.py, priority-samples
+a fixed number of points (2048 or 4096), normalizes, and writes one .npz per
+scene (~50KB at 2048, ~100KB at 4096).
+
+A fixed seed is used so every scene gets one deterministic sample -- no
+per-epoch sampling augmentation, every epoch sees the same points.
+
+Usage:
+  python -m s23dr_2026_example.make_sampled_cache \\
+      --in-dir cache/full --out-dir cache/sampled_2048 --seq-len 2048
+  python -m s23dr_2026_example.make_sampled_cache \\
+      --in-dir cache/full --out-dir cache/sampled_4096 --seq-len 4096
+
+The 3:1 colmap:depth quota ratio is fixed: at seq_len=2048 that's
+colmap=1536/depth=512; at seq_len=4096 that's colmap=3072/depth=1024.
+"""
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+# Priority sampling (same logic as train.py)
+def _priority_sample(source, group_id, seq_len, colmap_quota, depth_quota):
+    def pick(src_id, quota):
+        base = source == src_id
+        picked, remaining = [], quota
+        for tier in range(5):
+            if remaining <= 0:
+                break
+            pool = np.where(base & (group_id == tier))[0]
+            if len(pool) == 0:
+                continue
+            np.random.shuffle(pool)
+            take = min(remaining, len(pool))
+            picked.append(pool[:take])
+            remaining -= take
+        if remaining > 0:
+            pool = np.where(base & (group_id >= 0))[0]
+            if len(pool) > 0:
+                np.random.shuffle(pool)
+                picked.append(pool[:min(remaining, len(pool))])
+                remaining -= min(remaining, len(pool))
+        return np.concatenate(picked) if picked else np.array([], dtype=np.int64), remaining
+
+    idx_c, rem_c = pick(0, colmap_quota)
+    idx_d, rem_d = pick(1, depth_quota)
+
+    if rem_c > 0:
+        extra = np.setdiff1d(np.where((source == 1) & (group_id >= 0))[0], idx_d)
+        np.random.shuffle(extra)
+        idx_d = np.concatenate([idx_d, extra[:rem_c]])
+    if rem_d > 0:
+        extra = np.setdiff1d(np.where((source == 0) & (group_id >= 0))[0], idx_c)
+        np.random.shuffle(extra)
+        idx_c = np.concatenate([idx_c, extra[:rem_d]])
+
+    indices = np.concatenate([idx_c, idx_d])
+    num_valid = len(indices)
+    if num_valid < seq_len:
+        if num_valid == 0:
+            return np.zeros(seq_len, dtype=np.int64), np.zeros(seq_len, dtype=bool)
+        indices = np.concatenate([indices, np.full(seq_len - num_valid, indices[-1])])
+    mask = np.zeros(seq_len, dtype=bool)
+    mask[:num_valid] = True
+    return indices[:seq_len], mask
+
+
+def _process_sample(d, seq_len, colmap_q, depth_q):
+    """Sample and normalize one cached scene dict into a small npz-ready dict."""
+    xyz = np.asarray(d["xyz"], np.float32)
+    source = np.asarray(d["source"], np.uint8)
+    group_id = np.asarray(d["group_id"], np.int8)
+    class_id = np.asarray(d["class_id"], np.uint8)
+    vis_src = np.asarray(d["visible_src"], np.uint8)
+    vis_id = np.asarray(d["visible_id"], np.int16)
+    center = np.asarray(d["center"], np.float32)
+    scale = float(d["scale"])
+    gt_v = np.asarray(d["gt_vertices"], np.float32)
+    gt_e = np.asarray(d["gt_edges"], np.int32)
+
+    indices, mask = _priority_sample(source, group_id, seq_len, colmap_q, depth_q)
+    xyz_norm = ((xyz[indices] - center) / scale).astype(np.float32)
+    gt_seg = np.stack([gt_v[gt_e[:, 0]], gt_v[gt_e[:, 1]]], axis=1)
+    gt_seg_norm = ((gt_seg - center) / scale).astype(np.float32)
+
+    result = {
+        "xyz_norm": xyz_norm,
+        "class_id": class_id[indices].astype(np.uint8),
+        "source": source[indices].astype(np.uint8),
+        "mask": mask,
+        "gt_segments": gt_seg_norm,
+        "scale": np.float32(scale),
+        "center": center,
+        "gt_vertices": gt_v,
+        "gt_edges": gt_e,
+        "visible_src": vis_src[indices].astype(np.uint8),
+        "visible_id": vis_id[indices].astype(np.int16),
+    }
+    if "behind_gest_id" in d:
+        result["behind"] = np.asarray(d["behind_gest_id"], np.int16)[indices]
+    if "n_views_voted" in d:
+        result["n_views_voted"] = np.asarray(d["n_views_voted"], np.uint8)[indices]
+    if "vote_frac" in d:
+        result["vote_frac"] = np.asarray(d["vote_frac"], np.float32)[indices]
+    if "gt_edge_classes" in d:
+        result["gt_edge_classes"] = np.asarray(d["gt_edge_classes"], np.int64)
+    return result
+
+
+def main():
+    p = argparse.ArgumentParser(description="Stage 2: cached .pt -> sampled .npz")
+    p.add_argument("--in-dir", required=True, help="Directory of .pt files from cache_scenes.py")
+    p.add_argument("--out-dir", required=True, help="Output directory for .npz files")
+    p.add_argument("--seq-len", type=int, default=2048, help="Points per sample (2048 or 4096)")
+    p.add_argument("--seed", type=int, default=7)
+    args = p.parse_args()
+
+    colmap_q = args.seq_len * 3 // 4
+    depth_q = args.seq_len - colmap_q
+    print(f"seq_len={args.seq_len}  colmap={colmap_q}  depth={depth_q}")
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    np.random.seed(args.seed)
+
+    files = sorted(Path(args.in_dir).glob("*.pt"))
+    print(f"Found {len(files)} .pt files in {args.in_dir}")
+
+    done = 0
+    t0 = time.perf_counter()
+    for f in files:
+        out_f = out_dir / (f.stem + ".npz")
+        if out_f.exists():
+            done += 1
+            continue
+        d = torch.load(f, weights_only=False)
+        result = _process_sample(d, args.seq_len, colmap_q, depth_q)
+        np.savez(out_f, **result)
+        done += 1
+        if done % 2000 == 0:
+            rate = done / (time.perf_counter() - t0)
+            print(f"  {done}/{len(files)} [{rate:.0f}/s]")
+
+    elapsed = time.perf_counter() - t0
+    print(f"Done. {done} files in {elapsed:.0f}s -> {out_dir}")
+
+
+if __name__ == "__main__":
+    main()
+
+

s23dr_2026_example/model.py

@@ -0,0 +1,519 @@
+"""
+Perceiver-based transformer for 3D roof wireframe prediction.
+
+Architecture overview:
+
+    Input tokens [B, T, D]
+        |
+        v
+    input_proj: Linear -> GELU -> Linear -> LayerNorm   =>  [B, T, hidden]
+        |
+        v
+    Perceiver latent bottleneck (N PerceiverLatentLayers):
+        Learnable latent embeddings [L, hidden] are broadcast to batch.
+        Each layer: cross-attn(latents <- tokens) -> self-attn(latents) -> FFN
+        Output: latents [B, L, hidden]
+        |
+        v
+    Segment decoder (M SegmentDecoderLayers):
+        Learnable query embeddings [S, hidden] are broadcast to batch.
+        Each layer: cross-attn(queries <- latents) -> self-attn(queries) -> FFN
+        Output: queries [B, S, hidden]
+        |
+        v
+    segment_head: Linear -> 6D -> (midpoint, half_vector)
+        + query_offsets (learnable per-query bias)
+        endpoints = midpoint +/- half_vector  ->  [B, S, 2, 3]
+"""
+
+import torch
+import torch.nn as nn
+
+from .attention import MultiHeadSDPA, FeedForward
+
+
+# ---------------------------------------------------------------------------
+# Building blocks
+# ---------------------------------------------------------------------------
+
+class AttnResidual(nn.Module):
+    """Pre-norm attention + residual + dropout."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        kv_heads: int | None = None,
+        norm_class=None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        norm_class = norm_class or nn.LayerNorm
+        self.norm = norm_class(d_model)
+        self.attn = MultiHeadSDPA(d_model, num_heads, kv_heads=kv_heads, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        memory: torch.Tensor,
+        memory_key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        res = x
+        x = self.norm(x)
+        x = self.attn(x, memory, key_padding_mask=memory_key_padding_mask)
+        return res + self.drop(x)
+
+
+class FFNResidual(nn.Module):
+    """Pre-norm feed-forward + residual + dropout."""
+
+    def __init__(
+        self,
+        d_model: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        norm_class=None,
+    ):
+        super().__init__()
+        norm_class = norm_class or nn.LayerNorm
+        self.norm = norm_class(d_model)
+        self.ffn = FeedForward(d_model, dim_ff, activation=activation)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        res = x
+        x = self.norm(x)
+        x = self.ffn(x)
+        return res + self.drop(x)
+
+
+# ---------------------------------------------------------------------------
+# Perceiver encoder layer
+# ---------------------------------------------------------------------------
+
+class PerceiverLatentLayer(nn.Module):
+    """Single Perceiver latent layer.
+
+    If use_cross=True:  cross-attn(latents <- points) -> self-attn -> FFN
+    If use_cross=False: self-attn -> FFN  (saves compute in deep stacks)
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        kv_heads_cross: int | None = None,
+        kv_heads_self: int | None = None,
+        use_cross: bool = True,
+        norm_class=None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.use_cross = use_cross
+        if use_cross:
+            self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)
+
+    def forward(
+        self,
+        latents: torch.Tensor,
+        points: torch.Tensor,
+        points_key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.use_cross:
+            latents = self.cross(latents, points, memory_key_padding_mask=points_key_padding_mask)
+        latents = self.self_attn(latents, latents)
+        latents = self.ffn(latents)
+        return latents
+
+
+# ---------------------------------------------------------------------------
+# Segment decoder layer
+# ---------------------------------------------------------------------------
+
+class SegmentDecoderLayer(nn.Module):
+    """Single segment decoder layer.
+
+    cross-attn(queries <- latents) -> [cross-attn(queries <- inputs)] -> self-attn(queries) -> FFN
+
+    If input_xattn=True, adds a second cross-attention that attends directly
+    to the projected input tokens (bypassing the latent bottleneck). This gives
+    queries access to fine-grained point-level detail for vertex precision.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        kv_heads_cross: int | None = None,
+        kv_heads_self: int | None = None,
+        norm_class=None,
+        input_xattn: bool = False,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.input_xattn = input_xattn
+        if input_xattn:
+            self.cross_input = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)
+
+    def forward(
+        self,
+        queries: torch.Tensor,
+        latents: torch.Tensor,
+        src: torch.Tensor | None = None,
+        src_key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        queries = self.cross(queries, latents)
+        if self.input_xattn and src is not None:
+            queries = self.cross_input(queries, src, memory_key_padding_mask=src_key_padding_mask)
+        queries = self.self_attn(queries, queries)
+        queries = self.ffn(queries)
+        return queries
+
+
+# ---------------------------------------------------------------------------
+# Full model
+# ---------------------------------------------------------------------------
+
+class TokenTransformerSegments(nn.Module):
+    """Perceiver transformer that predicts 3D roof wireframe segments.
+
+    Takes point-cloud tokens and outputs segment endpoints as [B, S, 2, 3]
+    where S is the number of segments and each segment has two 3D endpoints.
+
+    Args:
+        segments:       Number of predicted segments (S).
+        in_dim:         Dimensionality of input tokens.
+        hidden:         Internal hidden dimension throughout the model.
+        num_heads:      Number of attention heads.
+        kv_heads_cross: Grouped-query heads for cross-attention (None = standard MHA).
+        kv_heads_self:  Grouped-query heads for self-attention (None = standard MHA).
+        dim_feedforward: FFN intermediate dimension.
+        dropout:        Dropout rate applied after attention and FFN.
+        latent_tokens:  Number of learnable latent embeddings (L) in the bottleneck.
+        latent_layers:  Number of PerceiverLatentLayers (N).
+        decoder_layers: Number of SegmentDecoderLayers (M).
+    """
+
+    def __init__(
+        self,
+        segments: int = 32,
+        in_dim: int = 128,
+        hidden: int = 128,
+        num_heads: int = 4,
+        kv_heads_cross: int | None = 2,
+        kv_heads_self: int | None = 0,
+        dim_feedforward: int = 256,
+        dropout: float = 0.01,
+        latent_tokens: int = 64,
+        latent_layers: int = 2,
+        decoder_layers: int = 2,
+        cross_attn_interval: int = 1,
+        norm_class=None,
+        activation: str = "gelu",
+        segment_conf: bool = False,
+        pre_encoder_layers: int = 0,
+        segment_param: str = "midpoint_halfvec",
+        length_floor: float = 0.0,
+        decoder_input_xattn: bool = False,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.segments = segments
+        self.out_vertices = segments * 2
+        self.segment_param = segment_param
+        self.decoder_input_xattn = decoder_input_xattn
+        norm_class = norm_class or nn.LayerNorm
+
+        # Treat 0 as "use standard MHA"
+        if kv_heads_cross is not None and kv_heads_cross <= 0:
+            kv_heads_cross = None
+        if kv_heads_self is not None and kv_heads_self <= 0:
+            kv_heads_self = None
+
+        # -- Input projection --
+        self.input_proj = nn.Sequential(
+            nn.Linear(in_dim, dim_feedforward),
+            nn.GELU(),
+            nn.Linear(dim_feedforward, hidden),
+            norm_class(hidden),
+        )
+
+        # -- Optional pre-encoder: self-attention on full token sequence --
+        if pre_encoder_layers > 0:
+            self.pre_encoder = nn.ModuleList([
+                SelfAttentionEncoderLayer(
+                    d_model=hidden,
+                    num_heads=num_heads,
+                    dim_ff=dim_feedforward,
+                    dropout=dropout,
+                    activation=activation,
+                    kv_heads=kv_heads_self,
+                    norm_class=norm_class,
+                    qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+                )
+                for _ in range(pre_encoder_layers)
+            ])
+        else:
+            self.pre_encoder = None
+
+        # -- Perceiver latent bottleneck --
+        self.latent_embed = nn.Embedding(latent_tokens, hidden)
+        N = latent_layers
+        self.latent_layers = nn.ModuleList([
+            PerceiverLatentLayer(
+                d_model=hidden,
+                num_heads=num_heads,
+                dim_ff=dim_feedforward,
+                dropout=dropout,
+                activation=activation,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                use_cross=(i == 0) or (i == N - 1) or (i % cross_attn_interval == 0),
+                norm_class=norm_class,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+            for i in range(N)
+        ])
+
+        # -- Segment decoder --
+        self.query_embed = nn.Embedding(segments, hidden)
+        self.decoder_layers = nn.ModuleList([
+            SegmentDecoderLayer(
+                d_model=hidden,
+                num_heads=num_heads,
+                dim_ff=dim_feedforward,
+                dropout=dropout,
+                activation=activation,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                norm_class=norm_class,
+                input_xattn=decoder_input_xattn,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+            for _ in range(decoder_layers)
+        ])
+
+        # -- Output head --
+        if segment_param == "midpoint_dir_len":
+            self.segment_head = nn.Linear(hidden, 7)  # mid(3) + dir(3) + len(1)
+        else:
+            self.segment_head = nn.Linear(hidden, 6)  # mid(3) + half(3)
+        self.query_offsets = nn.Parameter(torch.zeros(segments, 2, 3))
+
+        nn.init.trunc_normal_(self.segment_head.weight, mean=0.0, std=1e-3)
+        if self.segment_head.bias is not None:
+            nn.init.zeros_(self.segment_head.bias)
+            if segment_param == "midpoint_dir_len":
+                # softplus(0.5) * 0.1 ≈ 0.097 default length in normalized space
+                self.segment_head.bias.data[6] = 0.5
+        nn.init.normal_(self.query_offsets, mean=0.0, std=0.05)
+
+        # -- Optional confidence head --
+        self.segment_conf = segment_conf
+        if segment_conf:
+            self.conf_head = nn.Linear(hidden, 1)
+            nn.init.zeros_(self.conf_head.bias)
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor | list]:
+        """
+        Args:
+            tokens: Input point-cloud tokens [B, T, in_dim].
+            mask:   Boolean validity mask [B, T]. True = valid token.
+
+        Returns:
+            Dict with keys:
+                "vertices": [B, S*2, 3] flattened endpoints.
+                "segments": [B, S, 2, 3] segment endpoints.
+                "edges":    Per-batch list of (start, end) index pairs into vertices.
+                "conf":     [B, S] logits (only if segment_conf=True).
+        """
+        B = tokens.shape[0]
+
+        # Project input tokens
+        src = self.input_proj(tokens)  # [B, T, hidden]
+
+        # Padding mask (True where padded) for cross-attention
+        pad_mask = ~mask.bool() if mask is not None else None
+
+        # Optional pre-encoder: self-attention on full token sequence
+        if self.pre_encoder is not None:
+            for layer in self.pre_encoder:
+                src = layer(src, key_padding_mask=pad_mask)
+
+        # Perceiver latent bottleneck
+        latents = self.latent_embed.weight.unsqueeze(0).expand(B, -1, -1)
+        for layer in self.latent_layers:
+            latents = layer(latents, src, points_key_padding_mask=pad_mask)
+
+        # Segment decoder
+        queries = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1)
+        for layer in self.decoder_layers:
+            queries = layer(queries, latents,
+                            src=src if self.decoder_input_xattn else None,
+                            src_key_padding_mask=pad_mask if self.decoder_input_xattn else None)
+
+        # Predict segments -> endpoints
+        if self.segment_param == "midpoint_dir_len":
+            raw = self.segment_head(queries)  # [B, S, 7]
+            mid = raw[:, :, :3] + self.query_offsets[:, 0, :].unsqueeze(0)
+            direction = torch.nn.functional.normalize(raw[:, :, 3:6], dim=-1)
+            length = torch.nn.functional.softplus(raw[:, :, 6:7]) * 0.1
+            half = direction * length * 0.5
+        else:
+            raw = self.segment_head(queries).view(B, self.segments, 2, 3)
+            raw = raw + self.query_offsets.unsqueeze(0)
+            mid, half = raw[:, :, 0], raw[:, :, 1]
+        seg_params = torch.stack([mid - half, mid + half], dim=2)
+
+        vertices = seg_params.reshape(B, self.out_vertices, 3)
+        edges = [[(2 * i, 2 * i + 1) for i in range(self.segments)] for _ in range(B)]
+
+        out = {"vertices": vertices, "segments": seg_params, "edges": edges,
+               "src": src, "pad_mask": pad_mask, "queries": queries}
+        if self.segment_conf:
+            out["conf"] = self.conf_head(queries).squeeze(-1)  # [B, S]
+        return out
+
+
+# ---------------------------------------------------------------------------
+# Encoder-only layer (self-attention on full token sequence)
+# ---------------------------------------------------------------------------
+
+class SelfAttentionEncoderLayer(nn.Module):
+    """Single self-attention layer: self-attn(tokens) -> FFN."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        kv_heads: int | None = None,
+        norm_class=None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)
+
+    def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor | None = None) -> torch.Tensor:
+        x = self.self_attn(x, x, memory_key_padding_mask=key_padding_mask)
+        x = self.ffn(x)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# End-to-end model: tokenizer embeddings + perceiver
+# ---------------------------------------------------------------------------
+
+class EdgeDepthSegmentsModel(nn.Module):
+    """Tokenizer embeddings + transformer for 3D roof wireframes.
+
+    Supports two architectures via the `arch` parameter:
+    - "perceiver": Perceiver latent bottleneck (default, O(L*T) attention)
+    - "transformer": Standard self-attention encoder (O(T^2) attention)
+
+    Both share the same decoder, output head, and tokenizer.
+    """
+
+    def __init__(
+        self,
+        seq_cfg,
+        segments: int = 32,
+        hidden: int = 128,
+        num_heads: int = 4,
+        kv_heads_cross: int | None = 2,
+        kv_heads_self: int | None = 0,
+        dim_feedforward: int = 256,
+        dropout: float = 0.1,
+        latent_tokens: int = 64,
+        latent_layers: int = 1,
+        decoder_layers: int = 2,
+        label_emb_dim: int = 16,
+        src_emb_dim: int = 2,
+        behind_emb_dim: int = 8,
+        fourier_seed: int = 0,
+        cross_attn_interval: int = 1,
+        norm_class=None,
+        activation: str = "gelu",
+        segment_conf: bool = False,
+        use_vote_features: bool = False,
+        arch: str = "perceiver",
+        encoder_layers: int = 4,
+        pre_encoder_layers: int = 0,
+        segment_param: str = "midpoint_halfvec",
+        length_floor: float = 0.0,
+        decoder_input_xattn: bool = False,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+        learnable_fourier: bool = False,
+    ):
+        super().__init__()
+        self.seq_cfg = seq_cfg
+
+        from .tokenizer import EdgeDepthSequenceBuilder
+        self.tokenizer = EdgeDepthSequenceBuilder(
+            seq_cfg,
+            label_emb_dim=label_emb_dim,
+            src_emb_dim=src_emb_dim,
+            behind_emb_dim=behind_emb_dim,
+            fourier_seed=fourier_seed,
+            use_vote_features=use_vote_features,
+            learnable_fourier=learnable_fourier,
+        )
+
+        if arch == "transformer":
+            raise ValueError(
+                "arch='transformer' is no longer supported. "
+                "TransformerSegments has been removed; use arch='perceiver'.")
+        else:
+            self.segmenter = TokenTransformerSegments(
+                segments=segments,
+                in_dim=self.tokenizer.out_dim,
+                hidden=hidden,
+                num_heads=num_heads,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                dim_feedforward=dim_feedforward,
+                dropout=dropout,
+                latent_tokens=latent_tokens,
+                latent_layers=latent_layers,
+                decoder_layers=decoder_layers,
+                cross_attn_interval=cross_attn_interval,
+                norm_class=norm_class,
+                activation=activation,
+                segment_conf=segment_conf,
+                pre_encoder_layers=pre_encoder_layers,
+                segment_param=segment_param,
+                length_floor=length_floor,
+                decoder_input_xattn=decoder_input_xattn,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+
+    def forward_tokens(self, tokens: torch.Tensor, mask: torch.Tensor):
+        """Run the segmenter on pre-built token tensors."""
+        return self.segmenter(tokens, mask)

s23dr_2026_example/point_fusion.py

@@ -0,0 +1,554 @@
+"""
+point_fusion.py
+
+Simplified semantic point fusion for the 2026 dataset format.
+
+Takes per-view (ADE segmap, Gestalt segmap, depth) + sparse COLMAP point cloud
+from the usm3d/hoho22k_2026_trainval dataset and builds a compact, house-centric
+semantic point representation suitable for downstream wireframe prediction.
+
+Key differences from the 2025 pipeline:
+  - COLMAP is a ZIP of text files (cameras.txt, images.txt, points3D.txt)
+  - Depth is millimeter I;16 PNG (depth_scale=0.001 converts to meters)
+  - Views flagged with pose_only_in_colmap=True have zeroed K/R/t and must be
+    skipped for depth unprojection and projection
+  - Images arrive as PIL Images, not byte arrays
+"""
+
+from __future__ import annotations
+
+import zipfile
+from dataclasses import dataclass
+from io import BytesIO
+from typing import Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+from scipy.stats import mode as scipy_mode
+
+from .color_mappings import ade20k_color_mapping, gestalt_color_mapping
+
+# ---------------------------------------------------------------------------
+# Color packing helpers
+# ---------------------------------------------------------------------------
+
+def _pack_rgb_u32(rgb: np.ndarray) -> np.ndarray:
+    """Pack uint8 RGB (..., 3) into uint32 codes."""
+    rgb = rgb.astype(np.uint32, copy=False)
+    return (rgb[..., 0] << 16) | (rgb[..., 1] << 8) | rgb[..., 2]
+
+
+def _build_rgbcode_maps(color_mapping):
+    """Return (rgbcode_to_id, id_to_name) for a color mapping dict."""
+    names = list(color_mapping.keys())
+    rgbs = np.array([color_mapping[n] for n in names], dtype=np.uint8)
+    codes = _pack_rgb_u32(rgbs.reshape(-1, 1, 3)).reshape(-1)
+    rgbcode_to_id = {int(c): i for i, c in enumerate(codes)}
+    return rgbcode_to_id, names
+
+
+def _name_to_packed_rgb(name, mapping):
+    """Case-insensitive lookup returning a packed RGB code, or None."""
+    for key in mapping:
+        if key.lower() == name.lower():
+            rgb = np.array(mapping[key], np.uint8).reshape(1, 1, 3)
+            return int(_pack_rgb_u32(rgb).reshape(()))
+    return None
+
+# ---------------------------------------------------------------------------
+# Label mapping constants
+# ---------------------------------------------------------------------------
+
+ADE_RGBCODE_TO_ID, ADE_ID_TO_NAME = _build_rgbcode_maps(ade20k_color_mapping)
+GEST_RGBCODE_TO_ID, GEST_ID_TO_NAME = _build_rgbcode_maps(gestalt_color_mapping)
+NUM_ADE = len(ADE_ID_TO_NAME)
+NUM_GEST = len(GEST_ID_TO_NAME)
+
+GEST_INVALID_NAMES = ("unclassified", "unknown", "transition_line")
+GEST_INVALID_CODES = set(
+    int(_pack_rgb_u32(np.array(gestalt_color_mapping[n], np.uint8).reshape(1, 1, 3)).reshape(()))
+    for n in GEST_INVALID_NAMES if n in gestalt_color_mapping
+)
+
+# ADE classes whose surfaces are "see-through" for label fusion: when a point
+# projects onto one of these, we use the Gestalt label behind it instead.
+ADE_TRANSPARENT_NAMES = (
+    "wall", "building;edifice", "floor;flooring", "ceiling",
+    "windowpane;window", "door;double;door", "house", "skyscraper",
+    "screen;door;screen", "blind;screen", "hovel;hut;hutch;shack;shanty",
+    "tower", "booth;cubicle;stall;kiosk",
+)
+
+# ADE classes kept as "occluders/add-ons" when overlapping the house silhouette.
+ADE_OCCLUDER_ALLOWLIST_NAMES = (
+    "tree", "person;individual;someone;somebody;mortal;soul",
+    "car;auto;automobile;machine;motorcar", "truck;motortruck", "van",
+    "fence;fencing", "railing;rail",
+    "bannister;banister;balustrade;balusters;handrail",
+    "stairs;steps", "stairway;staircase", "step;stair", "pole",
+    "streetlight;street;lamp", "signboard;sign", "awning;sunshade;sunblind",
+    "plant;flora;plant;life", "pot;flowerpot",
+)
+
+# Precomputed arrays for the default name lists (avoids re-lookup every call).
+_DEFAULT_ADE_TRANSPARENT_CODES = np.array(
+    [c for n in ADE_TRANSPARENT_NAMES
+     if (c := _name_to_packed_rgb(n, ade20k_color_mapping)) is not None],
+    dtype=np.uint32,
+)
+_DEFAULT_ADE_OCCLUDER_IDS = np.array(
+    sorted({ADE_RGBCODE_TO_ID[c]
+            for n in ADE_OCCLUDER_ALLOWLIST_NAMES
+            if (c := _name_to_packed_rgb(n, ade20k_color_mapping)) is not None
+            and c in ADE_RGBCODE_TO_ID}),
+    dtype=np.int32,
+)
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+@dataclass(frozen=True)
+class FuserConfig:
+    """Simplified fusion configuration (no depth calibration fields)."""
+    depth_points_per_view: int = 20_000          # depth samples per view
+    depth_scale: float = 0.001                   # mm -> meters
+    depth_clip_percentile: float = 99.5          # drop extreme outliers
+    house_mask_dilate_px: int = 5                # dilate gestalt mask
+    min_support_views: int = 1                   # min views for a kept point
+    ade_transparent_classes: Tuple[str, ...] = ADE_TRANSPARENT_NAMES
+    ade_occluder_allowlist: Tuple[str, ...] = ADE_OCCLUDER_ALLOWLIST_NAMES
+
+# ---------------------------------------------------------------------------
+# Geometry: projection + depth unprojection
+# ---------------------------------------------------------------------------
+
+def project_world_points(points_world, K, R, t):
+    """Project (N,3) world points to pixel (u,v) with validity mask."""
+    pts = points_world.astype(np.float32, copy=False)
+    cam = (R @ pts.T + t).T  # (N, 3)
+    z = cam[:, 2]
+    valid = z > 1e-6
+    inv_z = np.zeros_like(z)
+    inv_z[valid] = 1.0 / z[valid]
+    x = cam[:, 0] * inv_z
+    y = cam[:, 1] * inv_z
+    u = K[0, 0] * x + K[0, 2]
+    v = K[1, 1] * y + K[1, 2]
+    return u, v, valid
+
+
+def unproject_depth_to_world(depth, K, R, t, num_points, sample_mask=None, rng=None):
+    """Convert a depth map + camera params to (M, 3) world points, M <= num_points."""
+    if rng is None:
+        rng = np.random.default_rng()
+    d = np.asarray(depth, dtype=np.float32)
+    if d.ndim != 2:
+        return np.zeros((0, 3), dtype=np.float32)
+
+    valid = np.isfinite(d) & (d > 1e-6)
+    if sample_mask is not None:
+        mask = np.asarray(sample_mask, dtype=bool)
+        if mask.shape != d.shape:
+            return np.zeros((0, 3), dtype=np.float32)
+        valid &= mask
+
+    ys, xs = np.where(valid)
+    if ys.size == 0:
+        return np.zeros((0, 3), dtype=np.float32)
+
+    idx = rng.choice(ys.size, size=min(num_points, ys.size), replace=False)
+    y = ys[idx].astype(np.float32)
+    x = xs[idx].astype(np.float32)
+    z = d[ys[idx], xs[idx]].astype(np.float32)
+
+    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
+    cam_pts = np.stack([(x - cx) * z / fx, (y - cy) * z / fy, z], axis=0)
+    # cam = R * world + t  =>  world = R^T * (cam - t)
+    world = (R.T @ (cam_pts - t)).T
+    return world.astype(np.float32, copy=False)
+
+
+def clean_depth(depth, clip_percentile):
+    """Clip extreme depth values."""
+    d = np.asarray(depth, dtype=np.float32)
+    d = np.where(np.isfinite(d), d, 0.0)
+    d[d <= 0] = 0.0
+    if clip_percentile is not None and clip_percentile > 0 and np.any(d > 0):
+        hi = float(np.percentile(d[d > 0], clip_percentile))
+        d = np.clip(d, 0.0, hi)
+    return d
+
+
+def dilate_mask(mask, radius_px):
+    """Binary dilation via cv2.  mask: (H, W) bool."""
+    if radius_px <= 0:
+        return mask
+    k = 2 * radius_px + 1
+    kernel = np.ones((k, k), np.uint8)
+    return cv2.dilate(mask.astype(np.uint8), kernel) > 0
+
+# ---------------------------------------------------------------------------
+# COLMAP extraction (2026 format)
+# ---------------------------------------------------------------------------
+
+def extract_colmap_points_2026(sample):
+    """Extract (N, 3) float32 COLMAP world points from a 2026-format sample.
+
+    sample['colmap'] must be a ZIP archive containing points3D.txt.
+    Fails fast if that file is missing (it is always present in the 2026 format).
+    """
+    colmap_blob = sample.get("colmap")
+    if colmap_blob is None:
+        return np.zeros((0, 3), dtype=np.float32)
+    if not isinstance(colmap_blob, (bytes, bytearray, memoryview)):
+        return np.zeros((0, 3), dtype=np.float32)
+
+    try:
+        with zipfile.ZipFile(BytesIO(colmap_blob)) as zf:
+            if "points3D.txt" not in set(zf.namelist()):
+                raise FileNotFoundError(
+                    "COLMAP ZIP is missing points3D.txt -- "
+                    "this is required in the 2026 dataset format")
+            with zf.open("points3D.txt") as f:
+                text = f.read().decode("utf-8", errors="ignore")
+            # Format: POINT3D_ID X Y Z R G B ERROR TRACK[]
+            # Filter comment/blank lines, parse columns 1-3 (X,Y,Z)
+            from io import StringIO
+            clean = "\n".join(l for l in text.split("\n") if l and not l.startswith("#"))
+            if not clean:
+                return np.zeros((0, 3), dtype=np.float32)
+            return np.loadtxt(StringIO(clean), dtype=np.float32, usecols=(1, 2, 3))
+    except zipfile.BadZipFile:
+        pass
+    return np.zeros((0, 3), dtype=np.float32)
+
+# ---------------------------------------------------------------------------
+# Label helpers
+# ---------------------------------------------------------------------------
+
+def _codes_from_image(img):
+    """Convert a PIL Image or numpy array to a (H, W) uint32 packed-RGB map."""
+    arr = np.asarray(img)
+    if arr.ndim == 2:
+        arr = np.stack([arr, arr, arr], axis=-1)
+    arr = arr[..., :3]
+    if arr.dtype != np.uint8:
+        arr = np.clip(arr, 0, 255).astype(np.uint8)
+    return _pack_rgb_u32(arr)
+
+
+def _row_majority(values):
+    """Row-wise majority vote on (P, V) int array; -1 means "no vote".
+    Returns (P,) with the most frequent non-negative value per row, or -1.
+
+    Masks -1 entries before voting so that abstentions don't outvote
+    actual labels (which happens when a point is visible in only 1-2 views).
+    """
+    P, V = values.shape
+    result = np.full(P, -1, dtype=values.dtype)
+
+    # For each row, find the most frequent non-negative value.
+    # Vectorized approach: flatten valid entries per row using argmax on counts.
+    # Since values are typically small non-negative ints (0-200), we can use
+    # a simple max-of-first-valid approach for speed when V is small.
+    for vi in range(V):
+        # For rows still unset, take the first valid vote
+        col = values[:, vi]
+        unset = result == -1
+        has_val = col >= 0
+        update = unset & has_val
+        result[update] = col[update]
+
+    # Now refine: if a row has multiple different valid votes, pick the mode.
+    # Check if any row has conflicting votes across views.
+    has_any = np.any(values >= 0, axis=1)
+    n_valid = np.sum(values >= 0, axis=1)
+    needs_vote = has_any & (n_valid > 1)
+
+    if np.any(needs_vote):
+        for i in np.where(needs_vote)[0]:
+            valid = values[i][values[i] >= 0]
+            # Use numpy bincount for speed (values are small non-neg ints)
+            counts = np.bincount(valid.astype(np.intp))
+            result[i] = counts.argmax()
+
+    return result
+
+# ---------------------------------------------------------------------------
+# Semantic fusion: house-centric, occluder-aware
+# ---------------------------------------------------------------------------
+
+def _fuse_labels_for_points(
+    points_world, Ks, Rs, ts, ade_images, gestalt_images,
+    ade_transparent_codes, ade_occluder_allowed_ids,
+    min_support_views, valid_view_mask=None,
+):
+    """Multi-view semantic label fusion with majority voting.
+
+    For each 3D point, project into every valid view:
+      - ADE "envelope" class -> use the Gestalt label behind it.
+      - ADE non-envelope -> keep if on the occluder allowlist.
+    Then majority-vote across views.
+
+    Returns dict: keep, visible_src, visible_id, behind_gest_id, support
+    """
+    P = points_world.shape[0]
+    V = min(len(Ks), len(Rs), len(ts), len(ade_images), len(gestalt_images))
+    empty = {
+        "keep": np.zeros(P, dtype=bool),
+        "visible_src": np.zeros(P, np.uint8),
+        "visible_id": np.full(P, -1, np.int16),
+        "behind_gest_id": np.full(P, -1, np.int16),
+        "support": np.zeros(P, np.uint8),
+    }
+    if P == 0 or V == 0:
+        return empty
+
+    # Per-view labels. src: 1=gestalt, 2=ade; -1 = no contribution.
+    visible_src_pv = np.full((P, V), -1, dtype=np.int8)
+    visible_id_pv = np.full((P, V), -1, dtype=np.int32)
+    behind_id_pv = np.full((P, V), -1, dtype=np.int32)
+    support = np.zeros(P, dtype=np.int32)
+
+    ade_allowed_set = set(ade_occluder_allowed_ids.tolist())
+    ade_transparent_u32 = ade_transparent_codes.astype(np.uint32, copy=False)
+    gest_invalid_arr = np.array(list(GEST_INVALID_CODES), dtype=np.uint32)
+
+    for vi in range(V):
+        if valid_view_mask is not None and not valid_view_mask[vi]:
+            continue
+
+        K = np.asarray(Ks[vi], np.float32)
+        R = np.asarray(Rs[vi], np.float32)
+        t = np.asarray(ts[vi], np.float32).reshape(3, 1)
+
+        ade_codes_img = _codes_from_image(ade_images[vi])
+        gest_codes_img = _codes_from_image(gestalt_images[vi])
+        H, W = ade_codes_img.shape
+
+        u, v, valid = project_world_points(points_world, K, R, t)
+        in_img = valid & (u >= 0) & (u < W) & (v >= 0) & (v < H)
+        if not np.any(in_img):
+            continue
+
+        ui = np.clip(np.round(u[in_img]).astype(np.int32), 0, W - 1)
+        vi_pix = np.clip(np.round(v[in_img]).astype(np.int32), 0, H - 1)
+        ade_codes = ade_codes_img[vi_pix, ui]
+        gest_codes = gest_codes_img[vi_pix, ui]
+
+        in_house = ~np.isin(gest_codes, gest_invalid_arr)
+        if not np.any(in_house):
+            continue
+
+        idx = np.where(in_img)[0][in_house]
+        ade_codes_h = ade_codes[in_house]
+        gest_codes_h = gest_codes[in_house]
+
+        behind_local = np.array(
+            [GEST_RGBCODE_TO_ID.get(int(c), -1) for c in gest_codes_h],
+            dtype=np.int32)
+        behind_id_pv[idx, vi] = behind_local
+
+        ade_is_transparent = np.isin(ade_codes_h, ade_transparent_u32)
+
+        # Case A: ADE is envelope -- use Gestalt label.
+        mask_a = ade_is_transparent & (behind_local >= 0)
+        if np.any(mask_a):
+            visible_src_pv[idx[mask_a], vi] = 1
+            visible_id_pv[idx[mask_a], vi] = behind_local[mask_a]
+
+        # Case B: ADE is non-envelope -- use ADE label (allowlist-filtered).
+        mask_b = ~ade_is_transparent
+        if np.any(mask_b):
+            ade_local = np.array(
+                [ADE_RGBCODE_TO_ID.get(int(c), -1) for c in ade_codes_h[mask_b]],
+                dtype=np.int32)
+            on_allowlist = np.array(
+                [int(a) in ade_allowed_set for a in ade_local], dtype=bool
+            ) & (ade_local >= 0)
+            if np.any(on_allowlist):
+                visible_src_pv[idx[mask_b][on_allowlist], vi] = 2
+                visible_id_pv[idx[mask_b][on_allowlist], vi] = ade_local[on_allowlist]
+
+        support[idx] += 1
+
+    # ---- Aggregate across views via majority vote ----
+    keep = (support >= min_support_views) & np.any(visible_src_pv >= 0, axis=1)
+
+    # Combine (src, id) into a single key for voting, then split back.
+    # src in {1,2} and id in [0, ~150], so stride=100k avoids collisions.
+    VIS_STRIDE = 100_000
+    vis_key = np.where(
+        visible_src_pv >= 0,
+        visible_src_pv.astype(np.int64) * VIS_STRIDE + visible_id_pv.astype(np.int64),
+        -1)
+    voted_key = _row_majority(vis_key)
+    voted_behind = _row_majority(behind_id_pv)
+
+    final_src = np.zeros(P, dtype=np.uint8)
+    final_id = np.full(P, -1, dtype=np.int16)
+    ok = voted_key >= 0
+    if np.any(ok):
+        final_src[ok] = (voted_key[ok] // VIS_STRIDE).astype(np.uint8)
+        final_id[ok] = (voted_key[ok] % VIS_STRIDE).astype(np.int16)
+
+    # ---- Vote confidence metadata ----
+    n_views_voted = np.sum(visible_src_pv >= 0, axis=1).astype(np.uint8)
+
+    # Fraction of voting views that agreed with the majority label
+    vote_frac = np.zeros(P, dtype=np.float32)
+    if np.any(ok):
+        for i in np.where(ok)[0]:
+            votes = vis_key[i][vis_key[i] >= 0]
+            if len(votes) > 0:
+                vote_frac[i] = (votes == voted_key[i]).sum() / len(votes)
+
+    return {
+        "keep": keep,
+        "visible_src": final_src,
+        "visible_id": final_id,
+        "behind_gest_id": voted_behind.astype(np.int16),
+        "support": support.astype(np.uint8),
+        "n_views_voted": n_views_voted,
+        "vote_frac": vote_frac,
+    }
+
+# ---------------------------------------------------------------------------
+# Compact scene builder (2026 dataset format)
+# ---------------------------------------------------------------------------
+
+def _resolve_ade_codes(cfg):
+    """Return (transparent_codes, occluder_ids) for the given config.
+    Uses precomputed module-level arrays when the config has default names.
+    """
+    if cfg.ade_transparent_classes == ADE_TRANSPARENT_NAMES:
+        transparent = _DEFAULT_ADE_TRANSPARENT_CODES
+    else:
+        transparent = np.array(
+            [c for n in cfg.ade_transparent_classes
+             if (c := _name_to_packed_rgb(n, ade20k_color_mapping)) is not None],
+            dtype=np.uint32)
+
+    if cfg.ade_occluder_allowlist == ADE_OCCLUDER_ALLOWLIST_NAMES:
+        occluder_ids = _DEFAULT_ADE_OCCLUDER_IDS
+    else:
+        occluder_ids = np.array(
+            sorted({ADE_RGBCODE_TO_ID[c]
+                    for n in cfg.ade_occluder_allowlist
+                    if (c := _name_to_packed_rgb(n, ade20k_color_mapping)) is not None
+                    and c in ADE_RGBCODE_TO_ID}),
+            dtype=np.int32)
+    return transparent, occluder_ids
+
+
+def _parse_gt_array(sample, key, dtype, expected_cols):
+    """Parse an optional ground-truth array from the sample dict."""
+    raw = sample.get(key)
+    if raw is None:
+        return None
+    arr = np.asarray(raw, dtype=dtype)
+    if arr.ndim == 2 and arr.shape[1] == expected_cols:
+        return arr
+    return None
+
+
+def build_compact_scene(sample, cfg, rng):
+    """Build a compact semantic point representation from a HuggingFace sample.
+
+    Expected sample keys: K, R, t, ade, gestalt, depth, colmap,
+    pose_only_in_colmap, wf_vertices (opt), wf_edges (opt), __key__ (opt).
+
+    Returns dict (xyz, source, visible_src, visible_id, behind_gest_id,
+    gt_vertices, gt_edges, sample_id) or None if no points survive fusion.
+    """
+    Ks = sample.get("K") or []
+    Rs = sample.get("R") or []
+    ts = sample.get("t") or []
+    ade_imgs = sample.get("ade") or []
+    gest_imgs = sample.get("gestalt") or []
+    depths = sample.get("depth") or []
+    pose_flags = sample.get("pose_only_in_colmap") or []
+
+    V = min(len(Ks), len(Rs), len(ts), len(ade_imgs), len(gest_imgs))
+    if V == 0:
+        return None
+
+    valid_view = [not (vi < len(pose_flags) and pose_flags[vi]) for vi in range(V)]
+    if not any(valid_view):
+        return None
+
+    # ---- COLMAP points ----
+    colmap_pts = extract_colmap_points_2026(sample)
+
+    # ---- Precompute house masks (from Gestalt), optionally dilated ----
+    gest_invalid_arr = np.array(list(GEST_INVALID_CODES), dtype=np.uint32)
+    house_masks = []
+    for vi in range(V):
+        if not valid_view[vi]:
+            house_masks.append(None)
+            continue
+        mask = ~np.isin(_codes_from_image(gest_imgs[vi]), gest_invalid_arr)
+        if cfg.house_mask_dilate_px > 0:
+            mask = dilate_mask(mask, cfg.house_mask_dilate_px)
+        house_masks.append(mask)
+
+    # ---- Sample depth points per view ----
+    depth_points_all = []
+    for vi in range(min(V, len(depths))):
+        if not valid_view[vi] or depths[vi] is None:
+            continue
+        d = clean_depth(
+            np.asarray(depths[vi], dtype=np.float32) * cfg.depth_scale,
+            cfg.depth_clip_percentile)
+        pts = unproject_depth_to_world(
+            depth=d,
+            K=np.asarray(Ks[vi], np.float32),
+            R=np.asarray(Rs[vi], np.float32),
+            t=np.asarray(ts[vi], np.float32).reshape(3, 1),
+            num_points=cfg.depth_points_per_view,
+            sample_mask=house_masks[vi], rng=rng)
+        if pts.shape[0]:
+            depth_points_all.append(pts)
+
+    # ---- Combine COLMAP + depth points ----
+    pts_list, src_list = [], []
+    if colmap_pts.shape[0]:
+        pts_list.append(colmap_pts)
+        src_list.append(np.zeros(colmap_pts.shape[0], dtype=np.uint8))   # 0=colmap
+    if depth_points_all:
+        all_depth = np.concatenate(depth_points_all, axis=0)
+        pts_list.append(all_depth)
+        src_list.append(np.ones(all_depth.shape[0], dtype=np.uint8))     # 1=depth
+    if not pts_list:
+        return None
+
+    points_world = np.concatenate(pts_list, axis=0).astype(np.float32, copy=False)
+    point_source = np.concatenate(src_list, axis=0).astype(np.uint8, copy=False)
+
+    # ---- Fuse semantic labels ----
+    ade_transparent_arr, ade_allow_ids = _resolve_ade_codes(cfg)
+    fused = _fuse_labels_for_points(
+        points_world=points_world, Ks=Ks, Rs=Rs, ts=ts,
+        ade_images=ade_imgs, gestalt_images=gest_imgs,
+        ade_transparent_codes=ade_transparent_arr,
+        ade_occluder_allowed_ids=ade_allow_ids,
+        min_support_views=cfg.min_support_views,
+        valid_view_mask=valid_view)
+
+    keep = fused["keep"]
+    if not np.any(keep):
+        return None
+
+    return {
+        "xyz": points_world[keep],
+        "source": point_source[keep],               # 0=colmap, 1=monodepth
+        "visible_src": fused["visible_src"][keep],   # 1=gestalt, 2=ade
+        "visible_id": fused["visible_id"][keep],
+        "behind_gest_id": fused["behind_gest_id"][keep],
+        "n_views_voted": fused["n_views_voted"][keep],
+        "vote_frac": fused["vote_frac"][keep],
+        "gt_vertices": _parse_gt_array(sample, "wf_vertices", np.float32, 3),
+        "gt_edges": _parse_gt_array(sample, "wf_edges", np.int64, 2),
+        "sample_id": sample.get("__key__", None),
+    }

s23dr_2026_example/postprocess_v2.py

@@ -0,0 +1,39 @@
+"""Post-processing functions for segment predictions."""
+import numpy as np
+
+
+def snap_to_point_cloud(vertices, xyz, class_id, snap_radius=0.5,
+                         target_classes=None):
+    """Snap vertices to nearby point cloud clusters of specific semantic classes."""
+    if target_classes is None:
+        target_classes = [1, 2]  # apex, eave_end_point
+
+    snapped = vertices.copy()
+    mask = np.isin(class_id, target_classes)
+
+    if mask.sum() < 2:
+        return snapped
+
+    target_pts = xyz[mask]
+
+    for i, v in enumerate(vertices):
+        dists = np.linalg.norm(target_pts - v, axis=-1)
+        close = dists < snap_radius
+        if close.sum() >= 2:
+            snapped[i] = target_pts[close].mean(axis=0)
+
+    return snapped
+
+
+def snap_horizontal(vertices, edges, max_slope=0.05):
+    """Snap near-horizontal edges to be exactly horizontal."""
+    verts = vertices.copy()
+    for a, b in edges:
+        a, b = int(a), int(b)
+        dy = abs(verts[a, 1] - verts[b, 1])
+        dxz = np.sqrt((verts[a, 0] - verts[b, 0])**2 + (verts[a, 2] - verts[b, 2])**2)
+        if dxz > 0.1 and dy / dxz < max_slope:
+            avg_y = 0.5 * (verts[a, 1] + verts[b, 1])
+            verts[a, 1] = avg_y
+            verts[b, 1] = avg_y
+    return verts

s23dr_2026_example/segment_postprocess.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def merge_vertices_iterative(vertices: np.ndarray, edges: np.ndarray,
+                              start: float = 0.15, end: float = 0.6,
+                              n_iters: int = 5):
+    """Iterative merge: start with tight threshold, gradually widen.
+
+    Avoids the worst transitive chaining effects of a single wide threshold.
+    Each pass merges only the closest pairs first, establishing stable cluster
+    centers before wider merges pull in more distant endpoints.
+
+    +0.004 HSS / +0.007 F1 over single-pass merge(0.4) on 1024 val samples.
+    """
+    pv, pe = vertices, edges
+    for t in np.linspace(start, end, n_iters):
+        pv, pe = merge_vertices(pv, pe, t)
+    return pv, pe
+
+
+def merge_vertices(vertices: np.ndarray, edges: np.ndarray, thresh: float):
+    verts = np.asarray(vertices, dtype=np.float32)
+    edges = np.asarray(edges, dtype=np.int64)
+    if verts.size == 0 or edges.size == 0:
+        return verts, edges
+
+    n = verts.shape[0]
+    parent = np.arange(n, dtype=np.int64)
+
+    def find(i):
+        while parent[i] != i:
+            parent[i] = parent[parent[i]]
+            i = parent[i]
+        return i
+
+    def union(i, j):
+        ri = find(i)
+        rj = find(j)
+        if ri != rj:
+            parent[rj] = ri
+
+    for i in range(n):
+        vi = verts[i]
+        for j in range(i + 1, n):
+            if np.linalg.norm(vi - verts[j]) <= thresh:
+                union(i, j)
+
+    clusters = {}
+    for i in range(n):
+        root = find(i)
+        clusters.setdefault(root, []).append(i)
+
+    new_vertices = []
+    mapping = {}
+    for new_idx, idxs in enumerate(clusters.values()):
+        pts = verts[idxs]
+        center = pts.mean(axis=0)
+        new_vertices.append(center)
+        for i in idxs:
+            mapping[i] = new_idx
+
+    new_edges = []
+    seen = set()
+    for a, b in edges:
+        na = mapping.get(int(a), int(a))
+        nb = mapping.get(int(b), int(b))
+        if na == nb:
+            continue
+        key = (na, nb) if na <= nb else (nb, na)
+        if key in seen:
+            continue
+        seen.add(key)
+        new_edges.append([na, nb])
+
+    return np.asarray(new_vertices, dtype=np.float32), np.asarray(new_edges, dtype=np.int64)

s23dr_2026_example/sinkhorn.py

@@ -0,0 +1,126 @@
+"""Sinkhorn optimal transport loss for segment matching.
+
+Note: at eps=0.05, sinkhorn gradients are near-zero (~1e-7 norm) for
+typical matrix sizes. The loss value is tracked but does not meaningfully
+train the model. Default sinkhorn_weight=0.0. See worklog.md for details.
+
+Future: schedule eps from large (1.0) to small (0.05) during training
+to get useful gradients early and precise matching late.
+"""
+import torch
+
+
+def batched_sinkhorn_loss(
+    pred_segments: torch.Tensor,
+    gt_pad: torch.Tensor,
+    gt_mask: torch.Tensor,
+    eps: float,
+    iters: int,
+    dustbin_cost: float | torch.Tensor,
+    pred_mass: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Batched sinkhorn segment matching loss.
+
+    Args:
+        pred_segments: [B, S, 2, 3] predicted segments
+        gt_pad:        [B, M, 2, 3] padded GT segments
+        gt_mask:       [B, M] bool mask (True = valid GT segment)
+        eps:           sinkhorn regularization
+        iters:         sinkhorn iterations
+        dustbin_cost:  cost for unmatched segments (scalar or [B])
+        pred_mass:     [B, S] per-segment mass weights (e.g. sigmoid(conf)).
+                       If None, uniform masses are used.
+
+    Returns:
+        [B] per-sample sinkhorn transport cost
+    """
+    B, S, _, _ = pred_segments.shape
+    M = gt_pad.shape[1]
+
+    # Allow per-sample dustbin cost
+    dc = torch.as_tensor(dustbin_cost, device=pred_segments.device, dtype=pred_segments.dtype)
+    if dc.dim() == 0:
+        dc = dc.expand(B)
+
+    # Compute cost matrices [B, S, M] in midpoint-halfvec space.
+    # Decouples position from direction: mid gradient is pure position,
+    # half gradient is pure direction/length. Sign-invariance on half
+    # handles segment direction ambiguity cleanly.
+    p0 = pred_segments[:, :, 0]  # [B, S, 3]
+    p1 = pred_segments[:, :, 1]  # [B, S, 3]
+    g0 = gt_pad[:, :, 0]         # [B, M, 3]
+    g1 = gt_pad[:, :, 1]         # [B, M, 3]
+
+    mid_pred = 0.5 * (p0 + p1)   # [B, S, 3]
+    half_pred = 0.5 * (p1 - p0)  # [B, S, 3]
+    mid_gt = 0.5 * (g0 + g1)     # [B, M, 3]
+    half_gt = 0.5 * (g1 - g0)    # [B, M, 3]
+
+    # Midpoint distance [B, S, M]
+    d_mid = torch.linalg.norm(
+        mid_pred.unsqueeze(2) - mid_gt.unsqueeze(1), dim=-1)
+
+    # Decoupled direction + length distance (sign-invariant for direction ambiguity)
+    len_pred = torch.linalg.norm(half_pred, dim=-1, keepdim=True).clamp(min=1e-6)  # [B, S, 1]
+    len_gt = torch.linalg.norm(half_gt, dim=-1, keepdim=True).clamp(min=1e-6)      # [B, M, 1]
+    dir_pred = half_pred / len_pred  # [B, S, 3]
+    dir_gt = half_gt / len_gt        # [B, M, 3]
+
+    # Direction distance: 1 - |cos(angle)|, sign-invariant [B, S, M]
+    cos_angle = (dir_pred.unsqueeze(2) * dir_gt.unsqueeze(1)).sum(dim=-1)  # [B, S, M]
+    d_dir = 1.0 - cos_angle.abs()
+
+    # Length distance [B, S, M]
+    d_len = (len_pred.unsqueeze(2) - len_gt.unsqueeze(1)).squeeze(-1).abs()
+
+    cost = d_mid + d_dir + d_len  # [B, S, M]
+
+    # Mask invalid GT segments with high cost so they go to dustbin
+    cost = torch.where(gt_mask.unsqueeze(1), cost, dc[:, None, None] * 10.0)
+
+    # Pad with dustbin row and column: [B, S+1, M+1]
+    cost_pad = dc[:, None, None].expand(B, S + 1, M + 1).clone()
+    cost_pad[:, :S, :M] = cost
+    cost_pad[:, -1, -1] = 0.0
+
+    # Masses
+    gt_counts = gt_mask.sum(dim=1).float()  # [B]
+
+    if pred_mass is not None:
+        # Confidence-weighted masses (matches learned_v2 approach).
+        # sigmoid(conf) gives per-segment mass; dustbin masses balance the totals.
+        # No normalization -- sum(a) == sum(b) == max(sum_pred, sum_gt).
+        pm = pred_mass.clamp(min=0.0)  # [B, S]
+        sum_pred = pm.sum(dim=1)       # [B]
+        sum_gt = gt_counts             # [B]
+        pred_dustbin = (sum_gt - sum_pred).clamp(min=0.0)  # [B]
+        gt_dustbin = (sum_pred - sum_gt).clamp(min=0.0)    # [B]
+        a = torch.cat([pm, pred_dustbin.unsqueeze(1)], dim=1)  # [B, S+1]
+        b_val = torch.zeros(B, M + 1, device=cost.device, dtype=cost.dtype)
+        b_val[:, :M] = gt_mask.float()  # 1.0 per valid GT segment
+        b_val[:, -1] = gt_dustbin
+    else:
+        # Uniform masses (normalized)
+        n = float(S)
+        denom = n + gt_counts  # [B]
+        a = (1.0 / denom).unsqueeze(1).expand(B, S + 1).clone()  # [B, S+1]
+        a[:, -1] = gt_counts / denom
+        b_val = (1.0 / denom).unsqueeze(1).expand(B, M + 1).clone()  # [B, M+1]
+        b_val[:, -1] = n / denom
+        # Zero out mass for invalid GT
+        b_val[:, :M] = b_val[:, :M] * gt_mask.float()
+
+    # Log-domain sinkhorn
+    log_a = torch.log(a + 1e-9)
+    log_b = torch.log(b_val + 1e-9)
+    log_k = -cost_pad / eps
+
+    log_u = torch.zeros_like(a)
+    log_v = torch.zeros_like(b_val)
+
+    for _ in range(iters):
+        log_u = log_a - torch.logsumexp(log_k + log_v.unsqueeze(1), dim=2)
+        log_v = log_b - torch.logsumexp(log_k + log_u.unsqueeze(2), dim=1)
+
+    transport = torch.exp(log_u.unsqueeze(2) + log_v.unsqueeze(1) + log_k)
+    return (transport * cost_pad).sum(dim=(1, 2))  # [B]

s23dr_2026_example/tokenizer.py

@@ -0,0 +1,88 @@
+"""Tokenizer: learned embeddings + Fourier features for the point cloud tokens.
+
+The EdgeDepthSequenceBuilder holds the learned embedding tables (label, source,
+behind) and the random Fourier positional encoding. At training time,
+build_tokens() in data.py applies these to pre-sampled point indices on GPU.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .point_fusion import NUM_ADE, NUM_GEST
+
+
+# -- Config --
+
+@dataclass(frozen=True)
+class EdgeDepthSequenceConfig:
+    seq_len: int = 2048
+    colmap_points: int = 1280
+    depth_points: int = 768
+    use_fourier: bool = True
+    fourier_dim: int = 32
+    fourier_scale: float = 10.0
+
+
+# -- Fourier positional encoding --
+
+class FourierFeatures(nn.Module):
+    def __init__(self, in_dim: int = 3, fourier_dim: int = 64,
+                 scale: float = 10.0, seed: int = 0,
+                 learnable: bool = False):
+        super().__init__()
+        gen = torch.Generator()
+        gen.manual_seed(seed)
+        B = torch.randn(fourier_dim, in_dim, generator=gen) * scale
+        if learnable:
+            self.B = nn.Parameter(B)
+        else:
+            self.register_buffer("B", B, persistent=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        proj = (2.0 * np.pi) * (x @ self.B.t())
+        return torch.cat([torch.sin(proj), torch.cos(proj)], dim=-1)
+
+
+# -- Sequence builder (holds embeddings) --
+
+class EdgeDepthSequenceBuilder(nn.Module):
+    """Holds learned embeddings for point cloud tokenization.
+
+    Used by the model at training time: build_tokens() calls
+    self.label_emb(class_id), self.src_emb(source), etc.
+    """
+
+    def __init__(self, cfg: EdgeDepthSequenceConfig, label_emb_dim: int = 16,
+                 src_emb_dim: int = 2, behind_emb_dim: int = 8,
+                 fourier_seed: int = 0, use_vote_features: bool = False,
+                 learnable_fourier: bool = False):
+        super().__init__()
+        self.cfg = cfg
+
+        self.num_labels = 13  # 11 structural + other_house + non_house
+        self.label_emb = nn.Embedding(self.num_labels, label_emb_dim)
+        self.src_emb = nn.Embedding(2, src_emb_dim)
+        self.behind_emb_dim = behind_emb_dim
+        if behind_emb_dim > 0:
+            self.behind_emb = nn.Embedding(NUM_GEST + 1, behind_emb_dim)
+
+        # Fourier positional encoding
+        if cfg.use_fourier:
+            self.pos_enc = FourierFeatures(
+                in_dim=3, fourier_dim=cfg.fourier_dim,
+                scale=cfg.fourier_scale, seed=fourier_seed,
+                learnable=learnable_fourier,
+            )
+            pos_dim = 3 + 2 * cfg.fourier_dim
+        else:
+            self.pos_enc = None
+            pos_dim = 3
+
+        vote_dim = 2 if use_vote_features else 0  # n_views_voted + vote_frac
+        self.use_vote_features = use_vote_features
+        self.out_dim = pos_dim + label_emb_dim + src_emb_dim + behind_emb_dim + vote_dim

s23dr_2026_example/train.py

@@ -0,0 +1,530 @@
+#!/usr/bin/env python3
+"""
+Training script for S23DR 2026.
+
+Usage:
+  python -m s23dr_2026_example.train --cache-dir hf://usm3d/s23dr-2026-sampled_2048_v2:train --steps 80000 --aug-rotate
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path as _Path
+if __package__ is None or __package__ == "":
+    _here = _Path(__file__).resolve().parent
+    if str(_here.parent) not in sys.path:
+        sys.path.insert(0, str(_here.parent))
+    __package__ = _here.name
+
+import argparse
+import gc
+import json
+import math
+import subprocess
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .tokenizer import EdgeDepthSequenceConfig
+from .model import EdgeDepthSegmentsModel
+from .data import build_loader, build_tokens
+from .losses import compute_loss, _loss_inner
+
+# Re-export for eval scripts
+from .data import HFCachedDataset, collate as _collate  # noqa: F401
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    p = argparse.ArgumentParser(description="S23DR 2026 training")
+    p.add_argument("--cache-dir", default=None, help="HF dataset path (hf://repo:split)")
+    p.add_argument("--val-cache-dir", default="", help="Separate cache for validation")
+    p.add_argument("--seq-len", type=int, default=2048,
+                   help="Input sequence length (2048 or 4096, must match dataset)")
+    p.add_argument("--arch", choices=["perceiver", "transformer"], default="perceiver",
+                   help="perceiver=latent bottleneck, transformer=full self-attention encoder")
+    p.add_argument("--segments", type=int, default=32)
+    p.add_argument("--hidden", type=int, default=128)
+    p.add_argument("--ff", type=int, default=512)
+    p.add_argument("--latent-tokens", type=int, default=128)
+    p.add_argument("--latent-layers", type=int, default=7)
+    p.add_argument("--encoder-layers", type=int, default=4,
+                   help="Encoder layers (transformer arch only)")
+    p.add_argument("--pre-encoder-layers", type=int, default=0,
+                   help="Self-attn layers on full token sequence before perceiver bottleneck")
+    p.add_argument("--decoder-layers", type=int, default=3)
+    p.add_argument("--decoder-input-xattn", action="store_true",
+                   help="Add cross-attention from segment queries to input tokens in each decoder layer")
+    p.add_argument("--qk-norm", action="store_true",
+                   help="Normalize Q and K per-head with learned temperature (stabilizes wide models)")
+    p.add_argument("--qk-norm-type", choices=["l2", "rms"], default="l2",
+                   help="QK-norm type: l2 (unit sphere) or rms (RMSNorm, preserves magnitudes)")
+    p.add_argument("--learnable-fourier", action="store_true",
+                   help="Make Fourier positional encoding learnable (vs fixed random)")
+    p.add_argument("--num-heads", type=int, default=4, help="Attention heads")
+    p.add_argument("--kv-heads-cross", type=int, default=2,
+                   help="KV heads for cross-attention (GQA; 0 = standard MHA)")
+    p.add_argument("--kv-heads-self", type=int, default=2,
+                   help="KV heads for self-attention (GQA; 0 = standard MHA)")
+    p.add_argument("--cross-attn-interval", type=int, default=4,
+                   help="Perceiver cross-attention frequency (every N latent layers)")
+    p.add_argument("--dropout", type=float, default=0.1)
+    p.add_argument("--weight-decay", type=float, default=0.01, help="AdamW weight decay")
+    p.add_argument("--steps", type=int, default=5000)
+    p.add_argument("--batch-size", type=int, default=32)
+    p.add_argument("--lr", type=float, default=3e-4)
+    p.add_argument("--adam-betas", default="0.9,0.95", help="AdamW beta1,beta2")
+    p.add_argument("--warmup", type=int, default=200, help="LR warmup steps")
+    p.add_argument("--cosine-decay", action="store_true",
+                   help="Cosine decay LR after warmup (to lr*0.01 at end)")
+    p.add_argument("--cooldown-start", type=int, default=0,
+                   help="Step to begin linear cooldown to lr*0.01 (0=disabled, constant LR after warmup)")
+    p.add_argument("--cooldown-steps", type=int, default=0,
+                   help="Number of steps for linear cooldown (0=no cooldown)")
+    p.add_argument("--seed", type=int, default=7)
+    p.add_argument("--deterministic", action="store_true",
+                   help="Force deterministic mode (disables torch.compile, slower but bit-reproducible)")
+    p.add_argument("--varifold-weight", type=float, default=0.0)
+    p.add_argument("--varifold-cross-only", action="store_true",
+                   help="Drop varifold self-energy (avoids O(S^2) spike, sinkhorn handles repulsion)")
+    p.add_argument("--sinkhorn-weight", type=float, default=1.0)
+    p.add_argument("--sinkhorn-eps", type=float, default=0.1,
+                   help="Sinkhorn regularization (larger = softer matching, stronger gradients)")
+    p.add_argument("--sinkhorn-eps-start", type=float, default=None,
+                   help="Starting eps for epsilon annealing (anneals to --sinkhorn-eps). None=no annealing.")
+    p.add_argument("--sinkhorn-eps-schedule", choices=["linear", "sqrt", "none"], default="none",
+                   help="Eps annealing schedule: linear, sqrt, or none (default: no annealing)")
+    p.add_argument("--sinkhorn-iters", type=int, default=20,
+                   help="Sinkhorn iterations")
+    p.add_argument("--sinkhorn-dustbin", type=float, default=0.3,
+                   help="Sinkhorn dustbin cost in normalized space")
+    p.add_argument("--endpoint-weight", type=float, default=0.0,
+                   help="Weight for endpoint distance loss (sinkhorn-matched, symmetric)")
+    p.add_argument("--endpoint-warmup", type=int, default=0,
+                   help="Steps to linearly warm up endpoint weight from 0 (0=instant)")
+    p.add_argument("--aug-rotate", action="store_true")
+    p.add_argument("--aug-jitter", type=float, default=0.0,
+                   help="Point position jitter std in normalized space (0=disabled, try 0.005)")
+    p.add_argument("--aug-drop", type=float, default=0.0,
+                   help="Fraction of points to randomly drop (0=disabled, try 0.1)")
+    p.add_argument("--aug-flip", action="store_true",
+                   help="Random mirror along X axis (50%% chance)")
+    p.add_argument("--rms-norm", action="store_true", default=True,
+                   help="Use RMSNorm (default). Use --no-rms-norm for LayerNorm")
+    p.add_argument("--no-rms-norm", dest="rms_norm", action="store_false")
+    p.add_argument("--activation", default="gelu", help="FFN activation: gelu, relu, relu_sq")
+    p.add_argument("--behind-emb-dim", type=int, default=8,
+                   help="Behind-gestalt embedding dim (0 to disable)")
+    p.add_argument("--vote-features", action="store_true",
+                   help="Add n_views_voted + vote_frac as raw token features (requires v2 data)")
+    p.add_argument("--segment-param", choices=["midpoint_halfvec", "midpoint_dir_len"],
+                   default="midpoint_halfvec",
+                   help="Output parameterization: halfvec (default) or decoupled direction+length")
+    p.add_argument("--length-floor", type=float, default=0.0,
+                   help="Minimum segment length for midpoint_dir_len (0=no floor)")
+    p.add_argument("--segment-conf", action="store_true",
+                   help="Add per-segment confidence head (use with --conf-thresh at eval)")
+    p.add_argument("--conf-weight", type=float, default=0.0,
+                   help="Weight for confidence loss (requires --segment-conf)")
+    p.add_argument("--conf-mode", choices=["sinkhorn", "sinkhorn_detach"], default="sinkhorn",
+                   help="Confidence training: 'match'=BCE, 'sinkhorn'=OT mass, 'sinkhorn_detach'=OT mass (detached)")
+    p.add_argument("--conf-clamp-min", type=float, default=None,
+                   help="Clamp conf logits to this minimum before sigmoid (e.g., -5)")
+    p.add_argument("--conf-head-wd", type=float, default=None,
+                   help="Separate weight decay for conf head (default: same as other params)")
+    p.add_argument("--ema-decay", type=float, default=0.0,
+                   help="EMA decay rate (0=disabled, try 0.9999). Saves EMA weights in checkpoints.")
+    p.add_argument("--out-dir", default=str(_Path(__file__).resolve().parent / "runs"))
+    p.add_argument("--resume", default="")
+    p.add_argument("--cpu", action="store_true")
+    p.add_argument("--args-from", default=None,
+                   help="Load defaults from a run's args.json (CLI flags override)")
+
+    # If --args-from is specified, load defaults from that JSON file first,
+    # then let CLI flags override.
+    raw_args = p.parse_args()
+    if raw_args.args_from is not None:
+        import json as _json
+        args_path = _Path(raw_args.args_from)
+        if not args_path.exists():
+            raise FileNotFoundError(f"--args-from file not found: {args_path}")
+        saved = _json.loads(args_path.read_text())
+        valid_dests = {a.dest for a in p._actions}
+        defaults = {}
+        for k, v in saved.items():
+            if k in valid_dests and k != "args_from":
+                defaults[k] = v
+        p.set_defaults(**defaults)
+        args = p.parse_args()
+        print(f"Loaded defaults from {args_path} (CLI flags override)")
+    else:
+        args = raw_args
+
+    # Validate required args
+    if not args.cache_dir:
+        p.error("--cache-dir is required (either directly or via --args-from)")
+
+    # Validate arg compatibility
+    if args.arch == "transformer":
+        perceiver_only = []
+        if args.latent_tokens != 128:
+            perceiver_only.append(f"--latent-tokens={args.latent_tokens}")
+        if args.latent_layers != 7:
+            perceiver_only.append(f"--latent-layers={args.latent_layers}")
+        if args.pre_encoder_layers != 0:
+            perceiver_only.append(f"--pre-encoder-layers={args.pre_encoder_layers}")
+        if args.cross_attn_interval != 4:
+            perceiver_only.append(f"--cross-attn-interval={args.cross_attn_interval}")
+        if perceiver_only:
+            raise ValueError(
+                f"Args {', '.join(perceiver_only)} have no effect with --arch transformer. "
+                f"Use --arch perceiver or remove them.")
+    if args.conf_weight > 0 and not args.segment_conf:
+        raise ValueError("--conf-weight requires --segment-conf")
+    if args.conf_mode in ("sinkhorn", "sinkhorn_detach") and args.sinkhorn_weight == 0:
+        raise ValueError("--conf-mode sinkhorn requires --sinkhorn-weight > 0")
+    if args.cosine_decay and args.cooldown_start > 0:
+        raise ValueError("--cosine-decay and --cooldown-start are mutually exclusive")
+
+    device = torch.device("cpu" if args.cpu else ("cuda" if torch.cuda.is_available() else "cpu"))
+    print(f"Device: {device}")
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    # Output
+    import hashlib, os
+    args_hash = hashlib.md5(json.dumps(vars(args), sort_keys=True).encode()).hexdigest()[:4]
+    run_tag = time.strftime("%Y%m%d_%H%M%S") + f"_{args_hash}_{os.getpid() % 10000:04d}"
+    out_dir = Path(args.out_dir) / run_tag
+    out_dir.mkdir(parents=True, exist_ok=True)
+    (out_dir / "checkpoints").mkdir(exist_ok=True)
+
+    # Tee stdout/stderr to run dir
+    import sys as _sys
+    _log_path = out_dir / "train.log"
+    class _Tee:
+        def __init__(self, path, stream):
+            self._file = open(path, "a")
+            self._stream = stream
+        def write(self, data):
+            self._stream.write(data)
+            self._file.write(data)
+            self._file.flush()
+        def flush(self):
+            self._stream.flush()
+            self._file.flush()
+    _sys.stdout = _Tee(_log_path, _sys.stdout)
+    _sys.stderr = _Tee(_log_path, _sys.stderr)
+
+    git_sha = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True,
+                             cwd=str(_Path(__file__).parent)).stdout.strip()
+    git_dirty = subprocess.run(["git", "diff", "--quiet"], capture_output=True,
+                               cwd=str(_Path(__file__).parent)).returncode != 0
+    run_info = {**vars(args), "git_sha": git_sha, "git_dirty": git_dirty}
+    (out_dir / "args.json").write_text(json.dumps(run_info, indent=2, sort_keys=True) + "\n")
+
+    # Set varifold cross-only mode before compile
+    if args.varifold_cross_only:
+        from . import losses as L
+        L.VARIFOLD_CROSS_ONLY = True
+        print("Varifold: cross-only mode (no self-energy)")
+
+    # Model
+    seq_len = args.seq_len
+    norm_class = torch.nn.RMSNorm if args.rms_norm else None
+    seq_cfg = EdgeDepthSequenceConfig(seq_len=seq_len)
+    model = EdgeDepthSegmentsModel(
+        seq_cfg=seq_cfg, segments=args.segments, hidden=args.hidden,
+        num_heads=args.num_heads, kv_heads_cross=args.kv_heads_cross,
+        kv_heads_self=args.kv_heads_self,
+        dim_feedforward=args.ff, dropout=args.dropout,
+        latent_tokens=args.latent_tokens, latent_layers=args.latent_layers,
+        decoder_layers=args.decoder_layers, cross_attn_interval=args.cross_attn_interval,
+        norm_class=norm_class, activation=args.activation,
+        segment_conf=args.segment_conf,
+        segment_param=args.segment_param,
+        length_floor=args.length_floor,
+        arch=args.arch, encoder_layers=args.encoder_layers,
+        pre_encoder_layers=args.pre_encoder_layers,
+        behind_emb_dim=args.behind_emb_dim,
+        use_vote_features=args.vote_features,
+        decoder_input_xattn=args.decoder_input_xattn,
+        qk_norm=args.qk_norm,
+        qk_norm_type=args.qk_norm_type,
+        learnable_fourier=args.learnable_fourier,
+    ).to(device)
+
+    try:
+        from torchinfo import summary
+        summary(model.segmenter,
+                input_data=[torch.zeros(1, seq_len, model.tokenizer.out_dim, device=device),
+                            torch.ones(1, seq_len, device=device, dtype=torch.bool)],
+                col_names=("input_size", "output_size", "num_params"), verbose=1)
+    except ImportError:
+        pass
+    print(f"Total params: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Compile (skip in deterministic mode for bit-reproducibility)
+    torch.set_float32_matmul_precision("high")
+    if args.deterministic:
+        torch.use_deterministic_algorithms(True)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        import os
+        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":16:8")
+        print("Deterministic mode: no torch.compile, bit-reproducible but ~3x slower")
+    elif device.type == "cuda":
+        model.segmenter = torch.compile(model.segmenter, mode="reduce-overhead", fullgraph=True)
+        from . import losses as L
+        L._loss_fn = torch.compile(_loss_inner, mode="reduce-overhead", fullgraph=True)
+        print("Compiled model + loss (reduce-overhead, fullgraph)")
+
+    # EMA
+    ema_model = None
+    if args.ema_decay > 0:
+        from copy import deepcopy
+        ema_model = deepcopy(model).eval()
+        for p_ema in ema_model.parameters():
+            p_ema.requires_grad_(False)
+        print(f"EMA enabled (decay={args.ema_decay})")
+
+    # Resume
+    start_step = 0
+    if args.resume:
+        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
+        try:
+            model.load_state_dict(ckpt["model"])
+        except RuntimeError:
+            state = {k.replace("segmenter._orig_mod.", "segmenter."): v
+                     for k, v in ckpt["model"].items()}
+            model.load_state_dict(state)
+        start_step = ckpt.get("step", 0)
+        print(f"Resumed from {args.resume} at step {start_step}")
+
+    betas = tuple(float(x) for x in args.adam_betas.split(","))
+
+    # Optimizer: AdamW with optional separate conf_head weight decay
+    conf_wd = args.conf_head_wd if args.conf_head_wd is not None else args.weight_decay
+    if args.conf_head_wd is not None:
+        conf_decay_params = []
+        other_params = []
+        for name, param in model.named_parameters():
+            if not param.requires_grad:
+                continue
+            if 'conf_head' in name:
+                conf_decay_params.append(param)
+            else:
+                other_params.append(param)
+        param_groups = [
+            {"params": other_params, "weight_decay": args.weight_decay},
+            {"params": conf_decay_params, "weight_decay": conf_wd},
+        ]
+        print(f"Conf head WD: {conf_wd} ({len(conf_decay_params)} params)")
+    else:
+        param_groups = model.parameters()
+
+    opt = torch.optim.AdamW(param_groups, lr=args.lr, weight_decay=args.weight_decay,
+                            betas=betas)
+    if args.resume and "optimizer" in ckpt:
+        opt.load_state_dict(ckpt["optimizer"])
+
+    # Data
+    torch.manual_seed(args.seed + 7919)
+    np.random.seed(args.seed + 7919)
+    train_loader = build_loader(args.cache_dir, args.batch_size, aug_rotate=args.aug_rotate,
+                                aug_jitter=args.aug_jitter, aug_drop=args.aug_drop,
+                                aug_flip=args.aug_flip)
+    val_loader = build_loader(args.val_cache_dir, args.batch_size) if args.val_cache_dir else None
+    data_iter = iter(train_loader)
+
+    # Intervals
+    log_int = max(1, min(50, args.steps // 20))
+    ckpt_int = 5000
+    val_int = ckpt_int if val_loader else 0
+
+    # Training loop
+    global_step = start_step
+    loss_ema, loss_sq_ema = 0.0, 0.0
+    t_start = time.perf_counter()
+
+    print(f"Training for {args.steps} steps | {args.segments}seg "
+          f"{args.hidden}h {args.latent_tokens}x{args.latent_layers}L "
+          f"{args.decoder_layers}D")
+
+    # Pre-fetch first batch
+    try:
+        next_batch = next(data_iter)
+    except StopIteration:
+        data_iter = iter(train_loader)
+        next_batch = next(data_iter)
+
+    # Freeze GC after setup to eliminate stalls during training
+    gc.collect()
+    gc.freeze()
+    gc.disable()
+
+    amp_ctx = torch.autocast(device_type='cuda', dtype=torch.bfloat16,
+                             enabled=(device.type == 'cuda'))
+
+    while global_step < args.steps:
+        tokens, masks, gt_list, scales, meta = build_tokens(next_batch, model, device)
+
+        # Epsilon annealing
+        if args.sinkhorn_eps_start is not None and args.sinkhorn_eps_start != args.sinkhorn_eps:
+            if args.sinkhorn_eps_schedule == "sqrt":
+                ratio_sq = (args.sinkhorn_eps_start / args.sinkhorn_eps) ** 2
+                t0 = max(args.steps * 0.8 / max(ratio_sq - 1, 1e-6), 1.0)
+                current_eps = args.sinkhorn_eps_start / math.sqrt(1 + global_step / t0)
+                current_eps = max(current_eps, args.sinkhorn_eps)
+            else:
+                frac = min(global_step / max(args.steps * 0.8, 1), 1.0)
+                current_eps = args.sinkhorn_eps_start + frac * (args.sinkhorn_eps - args.sinkhorn_eps_start)
+        else:
+            current_eps = args.sinkhorn_eps
+
+        with amp_ctx:
+            out = model.forward_tokens(tokens, masks)
+            pred = out["segments"]
+            conf = out.get("conf")
+
+            # Endpoint weight warmup
+            if args.endpoint_warmup > 0 and global_step < args.endpoint_warmup:
+                current_ep_w = args.endpoint_weight * global_step / args.endpoint_warmup
+            else:
+                current_ep_w = args.endpoint_weight
+
+            loss, terms = compute_loss(pred, gt_list, scales.to(device), device,
+                                       args.varifold_weight, args.sinkhorn_weight,
+                                       endpoint_w=current_ep_w,
+                                       conf_logits=conf, conf_weight=args.conf_weight,
+                                       conf_mode=args.conf_mode,
+                                       sinkhorn_eps=current_eps,
+                                       sinkhorn_iters=args.sinkhorn_iters,
+                                       sinkhorn_dustbin=args.sinkhorn_dustbin,
+                                       conf_clamp_min=args.conf_clamp_min)
+
+        loss_val = loss.item()
+        # Adaptive loss spike detection
+        if global_step < 100:
+            loss_ema = loss_val if global_step == start_step else 0.9 * loss_ema + 0.1 * loss_val
+            loss_sq_ema = loss_val**2 if global_step == start_step else 0.9 * loss_sq_ema + 0.1 * loss_val**2
+        else:
+            loss_ema = 0.99 * loss_ema + 0.01 * loss_val
+            loss_sq_ema = 0.99 * loss_sq_ema + 0.01 * loss_val**2
+        loss_std = max(math.sqrt(max(loss_sq_ema - loss_ema**2, 0)), 1e-6)
+        spike_thresh = loss_ema + 5 * loss_std
+
+        # Skip on total loss spike or NaN
+        if not math.isfinite(loss_val) or loss_val > max(spike_thresh, 0.5):
+            sample_ids = [m.get("sample_id", "?") for m in meta]
+            skip_reason = f"loss={loss_val:.2f} > thresh={spike_thresh:.2f}"
+            print(f"Step {global_step}: {skip_reason}, skipping (samples: {sample_ids[:3]})")
+            with open(out_dir / "skipped_samples.jsonl", "a") as f:
+                f.write(json.dumps({"step": global_step, "reason": skip_reason,
+                                    "samples": sample_ids}) + "\n")
+            try:
+                next_batch = next(data_iter)
+            except StopIteration:
+                data_iter = iter(train_loader)
+                next_batch = next(data_iter)
+            continue
+
+        opt.zero_grad()
+        loss.backward()
+
+        # Fetch next batch while GPU finishes backward
+        try:
+            next_batch = next(data_iter)
+        except StopIteration:
+            data_iter = iter(train_loader)
+            next_batch = next(data_iter)
+
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+
+        # LR schedule: warmup -> constant -> optional cooldown or cosine
+        if global_step < args.warmup:
+            lr = args.lr * (global_step + 1) / max(1, args.warmup)
+        elif args.cosine_decay:
+            progress = (global_step - args.warmup) / max(1, args.steps - args.warmup)
+            lr = args.lr * (0.01 + 0.99 * 0.5 * (1 + math.cos(math.pi * progress)))
+        elif args.cooldown_start > 0 and global_step >= args.cooldown_start:
+            progress = (global_step - args.cooldown_start) / max(1, args.cooldown_steps)
+            lr = args.lr * max(0.01, 1.0 - 0.99 * min(1.0, progress))
+        else:
+            lr = args.lr
+        for pg in opt.param_groups:
+            pg["lr"] = lr
+        opt.step()
+        global_step += 1
+
+        # EMA update
+        if ema_model is not None:
+            decay = args.ema_decay
+            with torch.no_grad():
+                for p_ema, p_model in zip(ema_model.parameters(), model.parameters()):
+                    p_ema.lerp_(p_model, 1.0 - decay)
+
+        # Log
+        entry = {"step": global_step, "ts": time.time(), "loss": loss.item(), "lr": lr}
+        entry.update({k: v.item() for k, v in terms.items()})
+        if global_step % log_int == 0:
+            grad_norm = sum(p.grad.norm().item()**2 for p in model.parameters()
+                            if p.grad is not None) ** 0.5
+            entry["grad_norm"] = grad_norm
+
+        if global_step % log_int == 0:
+            ms = (time.perf_counter() - t_start) / log_int * 1000
+            t_start = time.perf_counter()
+            t_str = " ".join(f"{k}={v:.4f}" for k, v in terms.items())
+            print(f"[{global_step}/{args.steps}] loss={loss.item():.4f} {t_str} "
+                  f"lr={lr:.2e} gnorm={entry.get('grad_norm', 0):.3f} [{ms:.0f}ms/step]")
+
+        if val_int > 0 and global_step % val_int == 0:
+            try:
+                vl_list = []
+                with torch.no_grad(), amp_ctx:
+                    for vb in val_loader:
+                        vt, vm, vg, vs, _ = build_tokens(vb, model, device)
+                        vo = model.forward_tokens(vt, vm)
+                        vl, _ = compute_loss(vo["segments"], vg, vs.to(device), device,
+                                             args.varifold_weight, args.sinkhorn_weight)
+                        if math.isfinite(vl.item()):
+                            vl_list.append(vl.item())
+                if vl_list:
+                    val_loss = float(np.mean(vl_list))
+                    print(f"  val_loss={val_loss:.4f}")
+                    entry["val_loss"] = val_loss
+            except Exception as e:
+                print(f"  val eval failed: {e}")
+
+        # Write log entry
+        with open(out_dir / "history.jsonl", "a") as f:
+            f.write(json.dumps(entry) + "\n")
+
+        if global_step % ckpt_int == 0:
+            try:
+                gc.enable(); gc.collect(); gc.freeze(); gc.disable()
+                torch.cuda.empty_cache()
+                save_dict = {"step": global_step, "model": model.state_dict(),
+                             "optimizer": opt.state_dict(), "args": vars(args)}
+                if ema_model is not None:
+                    save_dict["ema_model"] = ema_model.state_dict()
+                torch.save(save_dict, out_dir / "checkpoints" / f"step{global_step:06d}.pt")
+            except Exception as e:
+                print(f"  checkpoint save failed: {e}")
+
+    # Final save
+    save_dict = {"step": global_step, "model": model.state_dict(),
+                 "optimizer": opt.state_dict(), "args": vars(args)}
+    if ema_model is not None:
+        save_dict["ema_model"] = ema_model.state_dict()
+    torch.save(save_dict, out_dir / "checkpoints" / "final.pt")
+    print(f"Done. {global_step} steps. Output: {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

s23dr_2026_example/varifold.py

@@ -0,0 +1,53 @@
+import torch
+
+from .wire_varifold_kernels import (
+    loss_simpson3_batch,
+    loss_simpson3_mix_batch,
+)
+
+
+def segments_to_vertices_edges(segments: torch.Tensor):
+    segs = torch.as_tensor(segments, dtype=torch.float32)
+    vertices = segs.reshape(-1, 3)
+    edges = [(2 * i, 2 * i + 1) for i in range(segs.shape[0])]
+    return vertices, edges
+
+
+def varifold_loss_batch(
+    pred_segments: torch.Tensor,
+    gt_segments: torch.Tensor,
+    *,
+    sigma: float = 0.1,
+    variant: str = "semi_lobatto3",
+    t_nodes01: torch.Tensor | None = None,
+    t_w: torch.Tensor | None = None,
+    sigmas: torch.Tensor | None = None,
+    alpha: torch.Tensor | None = None,
+    normalize_alpha: bool = True,
+    len_pow: float | None = None,
+    gt_mask: torch.Tensor | None = None,
+    pred_weights: torch.Tensor | None = None,
+    cross_only: bool = False,
+) -> torch.Tensor:
+    if pred_segments.dim() != 4 or gt_segments.dim() != 4:
+        raise ValueError("pred_segments and gt_segments must be (B, N, 2, 3)")
+    p_pred, q_pred = pred_segments[:, :, 0], pred_segments[:, :, 1]
+    p_gt, q_gt = gt_segments[:, :, 0], gt_segments[:, :, 1]
+
+    w_gt = None
+    if gt_mask is not None:
+        w_gt = gt_mask.to(device=pred_segments.device, dtype=pred_segments.dtype)
+
+    w_pred = None
+    if pred_weights is not None:
+        w_pred = pred_weights.to(device=pred_segments.device, dtype=pred_segments.dtype)
+
+    if variant != "simpson3":
+        raise ValueError(
+            f"Unsupported varifold variant: {variant!r}. "
+            f"Only 'simpson3' is supported in batch mode.")
+    if sigmas is not None or alpha is not None:
+        if sigmas is None or alpha is None:
+            raise ValueError("sigmas and alpha are required for simpson3 mix")
+        return loss_simpson3_mix_batch(p_pred, q_pred, p_gt, q_gt, sigmas, alpha, w_gt=w_gt, w_pred=w_pred, normalize_alpha=normalize_alpha, cross_only=cross_only)
+    return loss_simpson3_batch(p_pred, q_pred, p_gt, q_gt, sigma, w_gt=w_gt, w_pred=w_pred)

s23dr_2026_example/wire_varifold_kernels.py

@@ -0,0 +1,168 @@
+import torch
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def segment_geom(p: torch.Tensor, q: torch.Tensor, eps: float = 1e-9):
+    """
+    p,q: (...,3)
+    returns d, a, ell, u:
+      d   = q - p
+      a   = ||d||^2
+      ell = sqrt(a + eps^2)
+      u   = d / ell
+    """
+    d = q - p
+    a = (d * d).sum(dim=-1)
+    eps_val = eps
+    if p.dtype in (torch.float16, torch.bfloat16):
+        eps_val = max(eps, float(torch.finfo(p.dtype).eps))
+    ell = torch.sqrt(a + eps_val * eps_val)
+    u = d / ell.unsqueeze(-1)
+    return d, a, ell, u
+
+def sample_points(p: torch.Tensor, q: torch.Tensor, nodes01: torch.Tensor):
+    # (...,3) + (K,) -> (...,K,3)
+    d = q - p
+    nodes = nodes01.to(device=p.device, dtype=p.dtype)
+    shape = [1] * (p.dim() - 1) + [nodes.shape[0], 1]
+    nodes = nodes.view(*shape)
+    return p.unsqueeze(-2) + nodes * d.unsqueeze(-2)
+
+
+# Fixed Lobatto-3 / Simpson nodes+weights on [0,1]
+LOBATTO3_NODES = torch.tensor([0.0, 0.5, 1.0])
+# LOBATTO3_W = torch.tensor([1.0/6.0, 4.0/6.0, 1.0/6.0])
+LOBATTO3_W = torch.tensor([1/3, 1/3, 1/3])
+LOBATTO3_W2 = LOBATTO3_W[:, None] * LOBATTO3_W[None, :]  # (3,3)
+
+
+def _prepare_mix_weights(sigmas, alpha, device, dtype, normalize_alpha: bool):
+    sigmas_t = torch.as_tensor(sigmas, device=device, dtype=dtype).clamp_min(1e-6)
+    alpha_t = torch.as_tensor(alpha, device=device, dtype=dtype)
+    if normalize_alpha:
+        alpha_t = alpha_t / alpha_t.sum().clamp_min(1e-12)
+    return sigmas_t, alpha_t
+
+# -----------------------------
+# Simpson-3 on both segments (3x3 product rule)
+# -----------------------------
+def _prep_weight(w, n: int, b: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None:
+    if w is None:
+        return None
+    w = torch.as_tensor(w, device=device, dtype=dtype)
+    if w.dim() == 1:
+        if w.shape[0] != n:
+            raise ValueError(f"weight length {w.shape[0]} != {n}")
+        w = w.unsqueeze(0).expand(b, -1)
+    elif w.dim() == 2:
+        if w.shape[0] != b or w.shape[1] != n:
+            raise ValueError(f"weight shape {tuple(w.shape)} != ({b}, {n})")
+    else:
+        raise ValueError("weights must be 1D or 2D")
+    return w
+
+
+def cross_simpson3(
+    pA,
+    qA,
+    pB,
+    qB,
+    sigma: float | torch.Tensor,
+    wA: torch.Tensor | None = None,
+    wB: torch.Tensor | None = None,
+):
+    device, dtype = pA.device, pA.dtype
+    batched = pA.dim() == 3
+    if not batched:
+        pA = pA.unsqueeze(0)
+        qA = qA.unsqueeze(0)
+        pB = pB.unsqueeze(0)
+        qB = qB.unsqueeze(0)
+    nodes = LOBATTO3_NODES.to(device=device, dtype=dtype)
+    w2 = LOBATTO3_W2.to(device=device, dtype=dtype)
+
+    bsz, nA, _ = pA.shape
+    nB = pB.shape[1]
+    wA = _prep_weight(wA, nA, bsz, device, dtype)
+    wB = _prep_weight(wB, nB, bsz, device, dtype)
+
+    _, _, ellA, uA = segment_geom(pA, qA)
+    _, _, ellB, uB = segment_geom(pB, qB)
+
+    XA = sample_points(pA, qA, nodes)  # (B,N,3,3)
+    YB = sample_points(pB, qB, nodes)  # (B,M,3,3)
+
+    # angular + length factors: (N,M)
+    ang = torch.matmul(uA, uB.transpose(-1, -2)).pow(2)
+    lenfac = ellA[:, :, None] * ellB[:, None, :]
+    if wA is not None or wB is not None:
+        if wA is None:
+            wA = torch.ones((bsz, nA), device=device, dtype=dtype)
+        if wB is None:
+            wB = torch.ones((bsz, nB), device=device, dtype=dtype)
+        lenfac = lenfac * (wA[:, :, None] * wB[:, None, :])
+
+    # spatial: build (N,M,3,3) kernel via broadcasting
+    diff = XA[:, :, None, :, None, :] - YB[:, None, :, None, :, :]  # (B,N,M,3,3,3)
+    r2 = (diff * diff).sum(dim=-1)                                 # (B,N,M,3,3)
+    sigma_t = torch.as_tensor(sigma, device=device, dtype=dtype)
+    if sigma_t.ndim == 0:
+        inv2s2 = 1.0 / (2.0 * sigma_t * sigma_t)
+    else:
+        if sigma_t.shape[0] != bsz:
+            raise ValueError(f"sigma batch {sigma_t.shape[0]} != {bsz}")
+        inv2s2 = (1.0 / (2.0 * sigma_t * sigma_t)).view(bsz, 1, 1, 1, 1)
+    K = torch.exp(-r2 * inv2s2)                                     # (B,N,M,3,3)
+
+    spatial = (K * w2).sum(dim=-1).sum(dim=-1)                     # (B,N,M)
+    out = (ang * lenfac * spatial).sum(dim=-1).sum(dim=-1)         # (B,)
+    return out[0] if not batched else out
+
+
+# -----------------------------
+# Batch losses
+# -----------------------------
+
+def loss_simpson3_batch(
+    p_pred: torch.Tensor,
+    q_pred: torch.Tensor,
+    p_gt: torch.Tensor,
+    q_gt: torch.Tensor,
+    sigma: float | torch.Tensor,
+    w_gt: torch.Tensor | None = None,
+    w_pred: torch.Tensor | None = None,
+    cross_only: bool = False,
+) -> torch.Tensor:
+    cross = cross_simpson3(p_pred, q_pred, p_gt, q_gt, sigma, wA=w_pred, wB=w_gt)
+    if cross_only:
+        # No self-energy: avoids O(S^2) blowup, sinkhorn handles repulsion
+        return -2.0 * cross
+    s_pred = cross_simpson3(p_pred, q_pred, p_pred, q_pred, sigma, wA=w_pred, wB=w_pred)
+    return s_pred - 2.0 * cross
+
+
+def loss_simpson3_mix_batch(
+    p_pred: torch.Tensor,
+    q_pred: torch.Tensor,
+    p_gt: torch.Tensor,
+    q_gt: torch.Tensor,
+    sigmas,
+    alpha,
+    w_gt: torch.Tensor | None = None,
+    w_pred: torch.Tensor | None = None,
+    normalize_alpha: bool = True,
+    cross_only: bool = False,
+) -> torch.Tensor:
+    device, dtype = p_pred.device, p_pred.dtype
+    sigmas_t = torch.as_tensor(sigmas, device=device, dtype=dtype).clamp_min(1e-6)
+    alpha_t = torch.as_tensor(alpha, device=device, dtype=dtype)
+    if normalize_alpha:
+        alpha_t = alpha_t / alpha_t.sum().clamp_min(1e-12)
+    if sigmas_t.ndim == 1:
+        losses = [loss_simpson3_batch(p_pred, q_pred, p_gt, q_gt, s, w_gt=w_gt, w_pred=w_pred, cross_only=cross_only) for s in sigmas_t]
+        return (torch.stack(losses, dim=0) * alpha_t[:, None]).sum(dim=0)
+    if sigmas_t.ndim == 2:
+        losses = [loss_simpson3_batch(p_pred, q_pred, p_gt, q_gt, sigmas_t[:, i], w_gt=w_gt, w_pred=w_pred, cross_only=cross_only) for i in range(sigmas_t.shape[1])]
+        return (torch.stack(losses, dim=0) * alpha_t[:, None]).sum(dim=0)
+    raise ValueError("sigmas must be 1D or 2D for batch loss")

script.py

@@ -1,131 +1,389 @@
-"""S23DR 2026 submission — sklearn edges + edge validation + improved heuristic."""
+"""S23DR 2026 submission: learned wireframe prediction from fused point clouds.
 
+Pipeline: raw sample -> point fusion -> priority sample 2048 -> model -> post-process -> wireframe
+"""
 from pathlib import Path
 from tqdm import tqdm
-import numpy as np
-from datasets import load_dataset
-from typing import Dict
-import os
 import json
-import gc
-import subprocess
+import os
 import sys
+import time
 
-# Ensure local imports work regardless of how script.py is invoked
-current_dir = str(Path(__file__).parent.absolute())
-if current_dir not in sys.path:
-    sys.path.insert(0, current_dir)
+import numpy as np
+import torch
 
-def install_if_missing(package):
-    try:
-        __import__(package)
-    except ImportError:
-        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
 
+def empty_solution():
+    return np.zeros((2, 3)), [(0, 1)]
 
-class Sample(Dict):
-    def __repr__(self):
-        return str({k: v.shape if hasattr(v, 'shape') else [type(v[0])] if isinstance(v, list) else type(v) for k, v in self.items()})
 
+# ---------------------------------------------------------------------------
+# Point fusion + sampling (from cache_scenes.py / make_sampled_cache.py)
+# ---------------------------------------------------------------------------
 
-def empty_solution():
-    return np.zeros((2, 3)), [(0, 1)]
+# Add our package to path
+SCRIPT_DIR = Path(__file__).resolve().parent
+sys.path.insert(0, str(SCRIPT_DIR))
 
+from s23dr_2026_example.point_fusion import build_compact_scene, FuserConfig
+from s23dr_2026_example.cache_scenes import (
+    _compute_group_and_class, _compute_smart_center_scale,
+)
+from s23dr_2026_example.make_sampled_cache import _priority_sample
 
-def process_sample(sample, i, sklearn_model=None):
+# Tokenizer / model imports
+from s23dr_2026_example.tokenizer import EdgeDepthSequenceConfig
+from s23dr_2026_example.model import EdgeDepthSegmentsModel
+from s23dr_2026_example.segment_postprocess import merge_vertices_iterative
+from s23dr_2026_example.varifold import segments_to_vertices_edges
+from s23dr_2026_example.postprocess_v2 import snap_to_point_cloud, snap_horizontal
+
+SEQ_LEN = 4096
+COLMAP_QUOTA = 3072
+DEPTH_QUOTA = 1024
+CONF_THRESH = 0.5
+MERGE_THRESH = 0.4
+SNAP_RADIUS = 0.5
+
+
+def fuse_and_sample(sample, cfg, rng):
+    """Run point fusion + priority sampling on a raw dataset sample.
+
+    Returns a dict with xyz_norm, class_id, source, mask, center, scale, etc.
+    ready for model inference. Returns None if fusion fails.
+    """
     try:
-        from sklearn_submission import predict_wireframe_sklearn
-        pred_vertices, pred_edges = predict_wireframe_sklearn(sample, sklearn_model)
+        scene = build_compact_scene(sample, cfg, rng)
     except Exception as e:
-        if i < 5:
-            print(f"  Sample {i} sklearn failed: {e}", flush=True)
-        try:
-            from hoho2025.example_solutions import predict_wireframe
-            pred_vertices, pred_edges = predict_wireframe(sample)
-        except Exception:
-            pred_vertices, pred_edges = empty_solution()
-    if i % 10 == 0:
-        gc.collect()
-    return {
-        'order_id': sample['order_id'],
-        'wf_vertices': np.array(pred_vertices).tolist(),
-        'wf_edges': [list(e) for e in pred_edges],
+        print(f"  Fusion failed: {e}")
+        return None
+
+    xyz = scene["xyz"]
+    source = scene["source"]
+
+    if len(xyz) < 10:
+        return None
+
+    # Compute group_id and class_id (same as cache_scenes.py)
+    behind_id = scene.get("behind_gest_id", np.full(len(xyz), -1, dtype=np.int16))
+    group_id, class_id = _compute_group_and_class(
+        scene["visible_src"], scene["visible_id"], behind_id, source)
+
+    # Normalize
+    center, scale = _compute_smart_center_scale(xyz, source)
+
+    # Priority sample
+    indices, mask = _priority_sample(source, group_id, SEQ_LEN, COLMAP_QUOTA, DEPTH_QUOTA)
+
+    xyz_norm = (xyz[indices] - center) / scale
+
+    result = {
+        "xyz_norm": xyz_norm.astype(np.float32),
+        "class_id": class_id[indices].astype(np.int64),
+        "source": source[indices].astype(np.int64),
+        "mask": mask,
+        "center": center.astype(np.float32),
+        "scale": np.float32(scale),
     }
 
+    # Optional fields
+    if "behind_gest_id" in scene:
+        behind = np.clip(scene["behind_gest_id"][indices].astype(np.int16), 0, None)
+        result["behind"] = behind.astype(np.int64)
+    if "n_views_voted" in scene:
+        result["n_views_voted"] = scene["n_views_voted"][indices].astype(np.float32)
+    if "vote_frac" in scene:
+        result["vote_frac"] = scene["vote_frac"][indices].astype(np.float32)
 
-if __name__ == "__main__":
-    print("------------ Loading dataset ------------", flush=True)
-    param_path = Path('params.json')
-    if param_path.exists():
-        with param_path.open() as f:
-            params = json.load(f)
+    # Visible src/id for snap post-processing
+    result["visible_src"] = scene["visible_src"][indices].astype(np.int64)
+    result["visible_id"] = scene["visible_id"][indices].astype(np.int64)
+
+    return result
+
+
+def load_model(checkpoint_path, device):
+    """Load model from checkpoint."""
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    args = ckpt.get("args", {})
+
+    norm_class = torch.nn.RMSNorm if args.get("rms_norm") else None
+    seq_cfg = EdgeDepthSequenceConfig(
+        seq_len=SEQ_LEN, colmap_points=COLMAP_QUOTA, depth_points=DEPTH_QUOTA)
+
+    model = EdgeDepthSegmentsModel(
+        seq_cfg=seq_cfg,
+        segments=args.get("segments", 64),
+        hidden=args.get("hidden", 256),
+        num_heads=args.get("num_heads", 4),
+        kv_heads_cross=args.get("kv_heads_cross", 2),
+        kv_heads_self=args.get("kv_heads_self", 2),
+        dim_feedforward=args.get("ff", 1024),
+        dropout=args.get("dropout", 0.1),
+        latent_tokens=args.get("latent_tokens", 256),
+        latent_layers=args.get("latent_layers", 7),
+        decoder_layers=args.get("decoder_layers", 3),
+        cross_attn_interval=args.get("cross_attn_interval", 4),
+        norm_class=norm_class,
+        activation=args.get("activation", "gelu"),
+        segment_conf=args.get("segment_conf", True),
+        behind_emb_dim=args.get("behind_emb_dim", 8),
+        use_vote_features=args.get("vote_features", True),
+        arch=args.get("arch", "perceiver"),
+        encoder_layers=args.get("encoder_layers", 4),
+        pre_encoder_layers=args.get("pre_encoder_layers", 0),
+        segment_param=args.get("segment_param", "midpoint_dir_len"),
+        qk_norm=args.get("qk_norm", True),
+    ).to(device)
+
+    # Handle torch.compile _orig_mod prefix
+    state = ckpt["model"]
+    fixed = {k.replace("segmenter._orig_mod.", "segmenter."): v
+             for k, v in state.items()}
+    model.load_state_dict(fixed, strict=True)
+    model.eval()
+    return model
+
+
+def build_tokens_single(sample_dict, model, device):
+    """Build token tensor for a single sample (no DataLoader)."""
+    xyz = torch.as_tensor(sample_dict["xyz_norm"], dtype=torch.float32).unsqueeze(0).to(device)
+    cid = torch.as_tensor(sample_dict["class_id"], dtype=torch.long).unsqueeze(0).to(device)
+    src = torch.as_tensor(sample_dict["source"], dtype=torch.long).unsqueeze(0).to(device)
+    masks = torch.as_tensor(sample_dict["mask"], dtype=torch.bool).unsqueeze(0).to(device)
+
+    B, T, _ = xyz.shape
+    tok = model.tokenizer
+    fourier = tok.pos_enc(xyz.reshape(-1, 3)).reshape(B, T, -1) \
+        if tok.pos_enc is not None else xyz.new_zeros(B, T, 0)
+    parts = [xyz, fourier, tok.label_emb(cid), tok.src_emb(src.clamp(0, 1))]
+
+    if tok.behind_emb_dim > 0:
+        if "behind" in sample_dict:
+            beh = torch.as_tensor(sample_dict["behind"], dtype=torch.long).unsqueeze(0).to(device)
+        else:
+            beh = xyz.new_zeros(B, T, dtype=torch.long)
+        parts.append(tok.behind_emb(beh))
+
+    if tok.use_vote_features:
+        if "n_views_voted" in sample_dict and "vote_frac" in sample_dict:
+            nv = ((torch.as_tensor(sample_dict["n_views_voted"], dtype=torch.float32).unsqueeze(0).to(device) - 2.7) / 1.0).unsqueeze(-1)
+            vf = ((torch.as_tensor(sample_dict["vote_frac"], dtype=torch.float32).unsqueeze(0).to(device) - 0.5) / 0.25).unsqueeze(-1)
+            parts.extend([nv, vf])
+        else:
+            parts.extend([xyz.new_zeros(B, T, 1), xyz.new_zeros(B, T, 1)])
+
+    tokens = torch.cat(parts, dim=-1)
+    return tokens, masks
+
+
+def predict_sample(sample_dict, model, device):
+    """Run model inference + post-processing on a fused sample.
+
+    Returns (vertices, edges) in world space.
+    """
+    tokens, masks = build_tokens_single(sample_dict, model, device)
+    scale = float(sample_dict["scale"])
+    center = sample_dict["center"]
+
+    with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.float16,
+                                          enabled=(device.type == 'cuda')):
+        out = model.forward_tokens(tokens, masks)
+
+    segs = out["segments"][0].float().cpu()
+    conf = torch.sigmoid(out["conf"][0].float()).cpu().numpy() if "conf" in out else None
+
+    # Confidence filter
+    if conf is not None:
+        keep = conf > CONF_THRESH
+        segs = segs[keep]
+    if len(segs) < 1:
+        return empty_solution()
+
+    # To world space
+    segs_world = segs.numpy() * scale + center
+
+    # Vertices + edges from segments
+    pv, pe = segments_to_vertices_edges(torch.tensor(segs_world))
+    pv, pe = pv.numpy(), np.array(pe, dtype=np.int32)
+
+    # Merge
+    pv, pe = merge_vertices_iterative(pv, pe)
+
+    # Snap to point cloud
+    xyz_norm = sample_dict["xyz_norm"]
+    mask = sample_dict["mask"]
+    cid = sample_dict["class_id"]
+    xyz_world = xyz_norm[mask] * scale + center
+    cid_valid = cid[mask]
+    pv = snap_to_point_cloud(pv, xyz_world, cid_valid, snap_radius=SNAP_RADIUS)
+
+    # Horizontal snap
+    pv = snap_horizontal(pv, pe)
+
+    if len(pv) < 2 or len(pe) < 1:
+        return empty_solution()
+
+    edges = [(int(a), int(b)) for a, b in pe]
+    return pv, edges
+
+def hybrid_merge(pred_v, pred_e, track_v, track_e, merge_radius=0.8):
+    if len(track_v) == 0:
+        return pred_v, pred_e
+        
+    pred_v = np.array(pred_v) if isinstance(pred_v, list) else pred_v
+    track_v = np.array(track_v)
+    
+    # We will append track vertices that are NOT close to any pred_v
+    if len(pred_v) > 0:
+        from scipy.spatial import cKDTree
+        tree = cKDTree(pred_v)
+        dists, indices = tree.query(track_v, k=1)
     else:
-        params = {"dataset": "usm3d/hoho22k_2026_trainval", "output_path": "."}
+        dists = np.full(len(track_v), np.inf)
+        indices = np.zeros(len(track_v), dtype=int)
+        
+    # Map track vertex indices to final vertex indices
+    track_to_final = {}
+    new_vertices = []
+    
+    for i, (d, idx) in enumerate(zip(dists, indices)):
+        if d <= merge_radius and len(pred_v) > 0:
+            # Map to existing pred_v
+            track_to_final[i] = int(idx)
+        else:
+            # Add as new vertex
+            track_to_final[i] = len(pred_v) + len(new_vertices)
+            new_vertices.append(track_v[i])
+            
+    final_v = list(pred_v) + new_vertices
+    final_e = list(pred_e)
+    
+    # Add track edges, mapping their indices
+    existing_edges = set()
+    for u, v in final_e:
+        existing_edges.add((min(u, v), max(u, v)))
+        
+    for u_t, v_t in track_e:
+        u_f = track_to_final.get(u_t)
+        v_f = track_to_final.get(v_t)
+        if u_f is not None and v_f is not None and u_f != v_f:
+            e = (min(u_f, v_f), max(u_f, v_f))
+            if e not in existing_edges:
+                final_e.append(e)
+                existing_edges.add(e)
+                
+    return np.array(final_v), final_e
 
-    data_path = Path('/tmp/data')
-    if not data_path.exists():
-        from huggingface_hub import snapshot_download
-        snapshot_download(repo_id=params['dataset'], local_dir="/tmp/data", repo_type="dataset")
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
 
-    os.system(f'ls -lahtrR {data_path}')
+if __name__ == "__main__":
+    t_start = time.time()
 
-    data_files = {}
-    public_tars = sorted([str(p) for p in data_path.rglob('*public*/**/*.tar')])
-    private_tars = sorted([str(p) for p in data_path.rglob('*private*/**/*.tar')])
-    if public_tars:
-        data_files["validation"] = public_tars
-    if private_tars:
-        data_files["test"] = private_tars
+    # Load params
+    param_path = Path("params.json")
+    with param_path.open() as f:
+        params = json.load(f)
+    print(f"Competition: {params.get('competition_id', '?')}")
+    print(f"Dataset: {params.get('dataset', '?')}")
 
-    loading_scripts = sorted(data_path.rglob('*.py'))
-    loading_script = str(loading_scripts[0]) if loading_scripts else str(data_path)
+    # Load test data
+    data_path = Path("/tmp/data")
+    if not data_path.exists():
+        from huggingface_hub import snapshot_download
+        snapshot_download(
+            repo_id=params["dataset"],
+            local_dir="/tmp/data",
+            repo_type="dataset",
+        )
 
+    from datasets import load_dataset
+    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 files: {data_files}")
     dataset = load_dataset(
-        loading_script, data_files=data_files,
-        trust_remote_code=True, writer_batch_size=100,
+        str(data_path / "hoho22k_2026_test_x_anon.py"),
+        data_files=data_files,
+        trust_remote_code=True,
+        writer_batch_size=100,
     )
-    print(f"Dataset: {dataset}", flush=True)
+    print(f"Loaded: {dataset}")
 
-    # Try to load sklearn model
-    sklearn_model = None
-    try:
-        install_if_missing('scikit-learn')
-        install_if_missing('networkx')
-        install_if_missing('scipy')
-        import pickle
-        model_path = Path(__file__).parent / 'sklearn_edge.pkl'
-        print(f"Looking for sklearn model at: {model_path} (exists={model_path.exists()})", flush=True)
-        if model_path.exists():
-            with open(model_path, 'rb') as f:
-                sklearn_model = pickle.load(f)
-            print("Loaded sklearn edge model OK", flush=True)
-        else:
-            print("sklearn model not found — using heuristic + edge validation only", flush=True)
-    except Exception as e:
-        print(f"sklearn failed: {e} — using heuristic + edge validation only", flush=True)
+    # Load model
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    checkpoint_path = SCRIPT_DIR / "checkpoint.pt"
+    
+    # Auto-download checkpoint if missing or just an LFS pointer
+    if not checkpoint_path.exists() or checkpoint_path.stat().st_size < 1000:
+        print("Downloading checkpoint.pt from upstream learned baseline...")
+        import urllib.request
+        ckpt_url = "https://huggingface.co/jacklangerman/s23dr-2026-submission/resolve/main/checkpoint.pt"
+        urllib.request.urlretrieve(ckpt_url, str(checkpoint_path))
+        print("Downloaded checkpoint.pt")
+        
+    model = load_model(checkpoint_path, device)
+    print(f"Model loaded: {sum(p.numel() for p in model.parameters()):,} params")
+
+    # Point fusion config
+    cfg = FuserConfig()
+    rng = np.random.RandomState(2718)
 
-    print("------------ Running predictions ---------------", flush=True)
+    # Process all samples
     solution = []
-    for subset_name in dataset.keys():
-        print(f"Predicting {subset_name}", flush=True)
-        for i, sample in enumerate(tqdm(dataset[subset_name])):
-            res = process_sample(sample, i, sklearn_model)
-            solution.append(res)
-            if i % 50 == 0:
-                print(f"  Processed {i} samples", flush=True)
-
-    print("------------ Saving results ---------------", flush=True)
-    output_path = Path(params.get('output_path', '.'))
-
-    with open(output_path / "submission.json", 'w') as f:
-        json.dump(solution, f)
-    print(f"Saved {len(solution)} predictions to submission.json", flush=True)
+    total_samples = sum(len(dataset[s]) for s in dataset)
+    processed = 0
 
-    try:
-        import pandas as pd
-        sub = pd.DataFrame(solution, columns=["order_id", "wf_vertices", "wf_edges"])
-        sub.to_parquet(output_path / "submission.parquet")
-    except Exception:
-        pass
+    for subset_name in dataset:
+        print(f"\nProcessing {subset_name} ({len(dataset[subset_name])} samples)...")
+
+        for sample in tqdm(dataset[subset_name], desc=subset_name):
+            order_id = sample["order_id"]
+
+            # Fuse + sample
+            fused = fuse_and_sample(sample, cfg, rng)
+            if fused is None:
+                pred_v, pred_e = empty_solution()
+            else:
+                try:
+                    pred_v, pred_e = predict_sample(fused, model, device)
+                    
+                    # Apply handcrafted triangulation tracking to catch missing corners/edges
+                    try:
+                        from triangulation import predict_wireframe_tracks
+                        # Force TRACK_MIN_VIEWS = 2 for aggressive recall
+                        track_v, track_e = predict_wireframe_tracks(sample, min_views=2)
+                        
+                        pred_v, pred_e = hybrid_merge(pred_v, pred_e, track_v, track_e, merge_radius=0.8)
+                    except Exception as track_e_err:
+                        print(f"  Track ensemble failed for {order_id}: {track_e_err}")
+                        
+                except Exception as e:
+                    print(f"  Predict failed for {order_id}: {e}")
+                    pred_v, pred_e = empty_solution()
+
+            solution.append({
+                "order_id": order_id,
+                "wf_vertices": pred_v.tolist() if isinstance(pred_v, np.ndarray) else pred_v,
+                "wf_edges": [(int(a), int(b)) for a, b in pred_e],
+            })
+            processed += 1
+
+            if processed % 50 == 0:
+                elapsed = time.time() - t_start
+                rate = elapsed / processed
+                remaining = (total_samples - processed) * rate
+                print(f"  [{processed}/{total_samples}] "
+                      f"{elapsed:.0f}s elapsed, ~{remaining:.0f}s remaining")
+
+    # Save
+    with open("submission.json", "w") as f:
+        json.dump(solution, f)
 
-    print("------------ Done ------------", flush=True)
+    elapsed = time.time() - t_start
+    print(f"\nDone. {processed} samples in {elapsed:.0f}s ({elapsed/max(processed,1):.1f}s/sample)")
+    print(f"Saved submission.json ({len(solution)} entries)")

sklearn_submission_diff.txt

@@ -0,0 +1,189 @@
+--- /Users/ihorivanyshyn/Documents/S23DR/s23dr-2026-submission/sklearn_submission.py	2026-04-26 12:52:14
++++ /Users/ihorivanyshyn/Documents/S23DR/handcrafted_submission_2026/sklearn_submission.py	2026-05-06 12:53:15
+@@ -84,7 +84,8 @@
+ # Stage 1 alone regressed in v16, but with DGCNN refinement the surviving
+ # candidates have median distance ~0.3 m to GT (vs ~1 m raw).
+ # v17 DGCNN vertex refinement — marginal on 100-sample sweep
+-# (ΔHSS +0.001 at best). Disabled by default.
++# (ΔHSS +0.001 at best). Disabled by default. Keep this conservative:
++# adding/removing vertices has a larger blast radius than adding edges.
+ USE_DGCNN_REFINEMENT = False
+ DGCNN_CLS_THRESHOLD = 0.5
+ DGCNN_DEDUP_RADIUS = 0.5
+@@ -100,7 +101,15 @@
+ #   t=0.7 +0.0039 (peak)  t=0.8 +0.0031
+ # Clean signal: F1 stable (±0.0006), IoU +0.0065 at t=0.7.
+ USE_DGCNN_EDGES = True
+-DGCNN_EDGE_THRESHOLD = 0.7
++# Ask the edge model for a wider candidate set, then apply our own
++# geometry gates below. This recovers medium-confidence true edges without
++# letting the classifier densify the graph unchecked.
++DGCNN_EDGE_THRESHOLD = 0.55
++DGCNN_EDGE_STRONG_THRESHOLD = 0.70
++DGCNN_EDGE_VERY_STRONG_THRESHOLD = 0.88
++DGCNN_EDGE_MAX_LENGTH = 8.0
++DGCNN_EDGE_MAX_PER_VERTEX = 2
++DGCNN_EDGE_REPROJ_DILATE_PX = 4
+ 
+ # v16: 3D vertex candidates from the S23DR 2025 winner Stage 1 — DISABLED.
+ # Raw cluster centroids without PointNet Stage 2 refinement have median
+@@ -191,7 +200,7 @@
+         device = "cuda" if _torch.cuda.is_available() else "cpu"
+     except Exception:
+         device = "cpu"
+-    _DGCNN_EDGE_MODEL = load_edge_model("checkpoints/edge_model_dgcnn.pt", device=device)
++    _DGCNN_EDGE_MODEL = load_edge_model("edge_model_dgcnn.pt", device=device)
+     return _DGCNN_EDGE_MODEL
+ 
+ 
+@@ -214,7 +223,7 @@
+         device = "cuda" if _torch.cuda.is_available() else "cpu"
+     except Exception:
+         device = "cpu"
+-    _DGCNN_VERTEX_MODEL = load_vertex_model("checkpoints/vertex_model_dgcnn.pt", device=device)
++    _DGCNN_VERTEX_MODEL = load_vertex_model("vertex_model_dgcnn.pt", device=device)
+     return _DGCNN_VERTEX_MODEL
+ 
+ # v7: ensemble with the standalone tracks-based predictor.
+@@ -500,8 +509,96 @@
+             if ok_views >= min_views:
+                 return True
+     return ok_views >= min_views
++
++
++def _passes_dgcnn_edge_gates(
++    v1: np.ndarray,
++    v2: np.ndarray,
++    prob: float,
++    all_xyz: np.ndarray,
++    kd_tree=None,
++    masks: dict | None = None,
++    views: dict | None = None,
++) -> bool:
++    """Conservative accept rule for learned edge candidates.
++
++    The DGCNN classifier is useful for recall, but raw learned edges can hurt
++    IoU if accepted without geometry. Strong candidates need COLMAP support;
++    very strong candidates may pass with looser sparse support; medium
++    candidates must also reproject onto gestalt edge pixels.
++    """
++    length = float(np.linalg.norm(v2 - v1))
++    if length < 0.25 or length > DGCNN_EDGE_MAX_LENGTH:
++        return False
++
++    strong_support = validate_edge(
++        v1, v2, all_xyz, kd_tree,
++        n_samples=24, radius=0.45, min_ratio=0.55,
++    )
++    if prob >= DGCNN_EDGE_STRONG_THRESHOLD and strong_support:
++        return True
++
++    loose_support = validate_edge(
++        v1, v2, all_xyz, kd_tree,
++        n_samples=24, radius=0.60, min_ratio=0.35,
++    )
++    if prob >= DGCNN_EDGE_VERY_STRONG_THRESHOLD and loose_support:
++        return True
++
++    if prob >= DGCNN_EDGE_STRONG_THRESHOLD and loose_support and masks and views:
++        return validate_edge_reprojection(
++            v1, v2, masks, views,
++            n_samples=24, min_views=1, min_hit_frac=0.35,
++        )
++
++    return False
++
++
++def _select_dgcnn_edges(
++    final_v: np.ndarray,
++    final_e: list,
++    dgcnn_edges: list,
++    all_xyz: np.ndarray,
++    kd_tree=None,
++    masks: dict | None = None,
++    views: dict | None = None,
++) -> list[tuple[int, int]]:
++    """Filter and degree-cap DGCNN edge proposals.
+ 
++    Existing edges are never removed here. At most
++    ``DGCNN_EDGE_MAX_PER_VERTEX`` learned edges are added at each vertex,
++    prioritising higher classifier probabilities.
++    """
++    existing = {tuple(sorted(e)) for e in final_e}
++    candidates = []
++    for i, j, prob in dgcnn_edges:
++        lo, hi = (int(i), int(j)) if i < j else (int(j), int(i))
++        if lo == hi or (lo, hi) in existing:
++            continue
++        prob = float(prob)
++        if _passes_dgcnn_edge_gates(
++            final_v[lo], final_v[hi], prob,
++            all_xyz, kd_tree, masks=masks, views=views,
++        ):
++            candidates.append((prob, lo, hi))
+ 
++    candidates.sort(reverse=True)
++    added_per_vertex = np.zeros(len(final_v), dtype=np.int32)
++    accepted: list[tuple[int, int]] = []
++    accepted_set = set()
++    for prob, lo, hi in candidates:
++        if (lo, hi) in accepted_set:
++            continue
++        if (added_per_vertex[lo] >= DGCNN_EDGE_MAX_PER_VERTEX
++                or added_per_vertex[hi] >= DGCNN_EDGE_MAX_PER_VERTEX):
++            continue
++        accepted.append((lo, hi))
++        accepted_set.add((lo, hi))
++        added_per_vertex[lo] += 1
++        added_per_vertex[hi] += 1
++    return accepted
++
++
+ def validate_edge(v1, v2, all_xyz, kd_tree=None, n_samples=20, radius=0.35, min_ratio=0.70):
+     """Check if edge v1→v2 is supported by COLMAP point cloud.
+ 
+@@ -1051,9 +1148,11 @@
+     if len(final_v) < 2 or len(final_e) < 1:
+         return empty_solution()
+ 
+-    # v18: DGCNN edge classifier — placed AFTER prune_not_connected so
+-    # that the vertex set is already fixed (no ghost vertices rescued by
+-    # spurious DGCNN edges). Only adds edges between surviving vertices.
++    # v19: guarded DGCNN edge rescue. The learned model is queried at a
++    # recall-friendly threshold, but new edges are accepted only if they
++    # also have sparse-cloud or reprojection evidence, then degree-capped.
++    # This targets the main weakness of v18: useful classifier recall
++    # without raw learned edges turning roofs into dense graphs.
+     if USE_DGCNN_EDGES and len(final_v) >= 2:
+         edge_model = _get_dgcnn_edge_model()
+         if edge_model is not None:
+@@ -1075,11 +1174,24 @@
+                         threshold=DGCNN_EDGE_THRESHOLD,
+                     )
+                     if dgcnn_edges:
+-                        existing = set(tuple(sorted(e)) for e in final_e)
+-                        for i, j, prob in dgcnn_edges:
+-                            lo, hi = (i, j) if i < j else (j, i)
+-                            if (lo, hi) not in existing:
+-                                final_e.append((lo, hi))
++                        masks, mvs_views = {}, {}
++                        try:
++                            masks, mvs_views = _build_gestalt_edge_masks(
++                                entry, dilate_px=DGCNN_EDGE_REPROJ_DILATE_PX,
++                            )
++                        except Exception:
++                            pass
++                        extra = _select_dgcnn_edges(
++                            np.asarray(final_v, dtype=np.float64),
++                            final_e,
++                            dgcnn_edges,
++                            all_xyz,
++                            kd_tree,
++                            masks=masks,
++                            views=mvs_views,
++                        )
++                        if extra:
++                            final_e.extend(extra)
+                 except Exception:
+                     pass
+ 

test_no_dgcnn.py

@@ -0,0 +1,32 @@
+import io
+import json
+import numpy as np
+from datasets import load_dataset
+from hoho2025.metric_helper import hss
+import sklearn_submission
+sklearn_submission.USE_DGCNN_EDGES = False
+sklearn_submission.USE_DGCNN_REFINEMENT = False
+
+# Turn on track ensemble parameters
+sklearn_submission.USE_TRACK_ENSEMBLE = True
+sklearn_submission.ADD_ISOLATED_TRACK_VERTICES = True
+
+# Turn on Tracks as Vertices
+sklearn_submission.USE_TRACKS_AS_VERTICES = True
+
+# Turn on winner candidates
+sklearn_submission.USE_WINNER_CANDIDATES = True
+
+print("Loading dataset...")
+dataset = load_dataset('usm3d/hoho22k_2026_trainval', split='train', streaming=True, trust_remote_code=True)
+samples = [s for idx, s in enumerate(dataset) if idx < 10]
+
+scores = []
+for idx, sample in enumerate(samples):
+    pred_v, pred_e = sklearn_submission.predict_wireframe_sklearn(sample)
+    gt_v = sample.get('wf_vertices')
+    gt_e = sample.get('wf_edges')
+    if gt_v is not None and gt_e is not None:
+        scores.append(hss(pred_v, pred_e, gt_v, gt_e).hss)
+
+print(f"Avg HSS without DGCNN: {sum(scores)/len(scores):.4f}")

test_triangulation.py

@@ -0,0 +1,20 @@
+import os
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
+import numpy as np
+from datasets import load_dataset
+import sklearn_submission
+from triangulation import predict_wireframe_tracks
+
+print("Loading dataset...")
+dataset = load_dataset('usm3d/hoho22k_2026_trainval', split='train', streaming=True, trust_remote_code=True)
+samples = [s for idx, s in enumerate(dataset) if idx < 3]
+
+for idx, sample in enumerate(samples):
+    print(f"Sample {idx}")
+    try:
+        tv, te = predict_wireframe_tracks(sample)
+        print(f"Tracks: {len(tv)} vertices, {len(te)} edges")
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        print(f"CRASH: {e}")

time_test.py

@@ -0,0 +1,29 @@
+import time
+import os
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
+import numpy as np
+from datasets import load_dataset
+import sklearn_submission
+from sklearn_submission import predict_wireframe_sklearn
+
+sklearn_submission.USE_BUNDLE_ADJUST = True
+sklearn_submission.ADD_ISOLATED_TRACK_VERTICES = True
+
+print("Loading dataset...")
+dataset = load_dataset('usm3d/hoho22k_2026_trainval', split='train', streaming=True, trust_remote_code=True)
+
+# Process 5 samples and time them
+times = []
+for idx, sample in enumerate(dataset):
+    if idx >= 5:
+        break
+    start = time.time()
+    try:
+        predict_wireframe_sklearn(sample)
+    except Exception as e:
+        print(f"Error on sample {idx}: {e}")
+    elapsed = time.time() - start
+    times.append(elapsed)
+    print(f"Sample {idx}: {elapsed:.2f} seconds")
+
+print(f"Average time per sample: {np.mean(times):.2f} seconds")

tune_dgcnn.py

@@ -0,0 +1,100 @@
+import io
+import json
+import numpy as np
+import sys
+import os
+import random
+import importlib
+from datasets import load_dataset
+from hoho2025.metric_helper import hss
+
+import sklearn_submission
+
+print("Loading dataset...")
+dataset = load_dataset('usm3d/hoho22k_2026_trainval', split='train', streaming=True, trust_remote_code=True)
+
+num_samples = 10
+samples = []
+for idx, s in enumerate(dataset):
+    if idx >= num_samples:
+        break
+    samples.append(s)
+
+def set_hyperparameters(p):
+    sklearn_submission.DGCNN_EDGE_THRESHOLD = p['threshold']
+    sklearn_submission.DGCNN_EDGE_STRONG_THRESHOLD = p['strong_threshold']
+    sklearn_submission.DGCNN_EDGE_VERY_STRONG_THRESHOLD = p['very_strong_threshold']
+    sklearn_submission.DGCNN_EDGE_MAX_LENGTH = p['max_length']
+    sklearn_submission.DGCNN_EDGE_MAX_PER_VERTEX = p['max_per_vertex']
+    sklearn_submission.DGCNN_EDGE_REPROJ_DILATE_PX = p['dilate_px']
+
+def evaluate():
+    scores = []
+    for sample in samples:
+        try:
+            pred_v, pred_e = sklearn_submission.predict_wireframe_sklearn(sample)
+        except Exception as e:
+            pred_v, pred_e = np.zeros((2, 3)), [(0, 1)]
+            
+        gt_v = sample.get('wf_vertices')
+        gt_e = sample.get('wf_edges')
+        
+        if gt_v is None or gt_e is None:
+            continue
+            
+        res = hss(pred_v, pred_e, gt_v, gt_e)
+        scores.append(res.hss)
+    
+    import gc
+    gc.collect()
+    
+    return sum(scores) / len(scores) if scores else 0
+
+space = {
+    'threshold': [0.45, 0.50, 0.55, 0.60],
+    'strong_threshold': [0.65, 0.70, 0.75],
+    'very_strong_threshold': [0.85, 0.88, 0.90, 0.95],
+    'max_length': [6.0, 8.0, 10.0],
+    'max_per_vertex': [1, 2, 3],
+    'dilate_px': [2, 4, 6]
+}
+
+trials = 8
+print(f"Running Random Search for {trials} trials over {num_samples} samples...")
+
+best_score = -1.0
+best_params = None
+
+# Base parameters evaluation
+base_p = {
+    'threshold': 0.55,
+    'strong_threshold': 0.70,
+    'very_strong_threshold': 0.88,
+    'max_length': 8.0,
+    'max_per_vertex': 2,
+    'dilate_px': 4
+}
+set_hyperparameters(base_p)
+base_score = evaluate()
+print(f"Baseline Score: {base_score:.4f} | Params: {base_p}")
+best_score = base_score
+best_params = base_p
+
+for i in range(trials):
+    p = {k: random.choice(v) for k, v in space.items()}
+    set_hyperparameters(p)
+    score = evaluate()
+    print(f"Trial {i+1} Score: {score:.4f} | Params: {p}")
+    if score > best_score:
+        best_score = score
+        best_params = p
+
+print(f"\n--- BEST RESULT ---")
+print(f"Best HSS Score: {best_score:.4f}")
+print("Best Parameters:")
+for k, v in best_params.items():
+    print(f"  {k}: {v}")
+
+# Write the best params to a JSON for record
+with open("best_dgcnn_params.json", "w") as f:
+    json.dump(best_params, f, indent=4)