Add DINOv2 edge classifier (v4) for post-processing edge filtering

Trained head: 128K params, frozen DINOv2-small (22M params) backbone.
Operates as a final post-processing step: for each predicted edge, sample
DINOv2 patch features at projected edge midpoint per view, mean+max pool
across views (768-dim), concat with 40-dim geometric+mask features (808-dim),
pass through 4-layer MLP head -> P(keep).

Best val acc 82.1% (vs 70-75% for hand-engineered features alone).
Operating point: thresh=0.15, min_keep=0.85 - drop only the bottom ~15% of
edges where the classifier is most confident they are wrong. Re-runs orphan
drop after to clean up any stranded vertices.

Local A/B vs production:
100 samples: +0.0032 hss_mean (t=+1.07, 46 wins / 44 losses)
200 samples: +0.0030 hss_mean (t=+1.26, 100 wins / 77 losses)

DINOv2 weights download automatically on first run via torch.hub.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

edge_classifier_v2.py

@@ -0,0 +1,339 @@
+"""V2 edge classifier with image-feature inputs.
+
+Richer features than v1: per-edge-class support along projection, endpoint
+distance to gestalt corner pixels, view-consistency stats. The signal is what
+the gestalt segmentation actually shows along the projected edge — not just
+"are there any edge pixels", but WHICH class and HOW consistently.
+
+Feature layout (40 total):
+    0-11   geometric features (same as v1)
+   12-18   per-edge-class mean fraction along projected segment (7 classes)
+   19-25   per-edge-class MAX fraction (best single view)
+   26      endpoint A: mean px distance to nearest gestalt corner (clipped 30)
+   27      endpoint B: same
+   28      endpoint A: count of views where it's <10px from a corner
+   29      endpoint B: same
+   30      edge midpoint: mean px dist to nearest gestalt edge pixel
+   31      num views with any-edge support > 0.3
+   32      num views with any-edge support > 0.6
+   33      mean projected-length in pixels (across views)
+   34      std of any-edge support across views (consistency)
+   35      mean depth z (camera frame, across views) — flat-roof prior
+   36      colmap support: min endpoint dist to colmap point
+   37      colmap support: midpoint dist to colmap point
+   38      edge length / scene median length (raw ratio)
+   39      vertex_count / scene_median_count (graph context)
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+EDGE_CLASSES = (
+    "eave", "ridge", "rake", "valley", "hip", "flashing", "step_flashing",
+)
+POINT_CLASSES = ("apex", "eave_end_point", "flashing_end_point")
+NUM_FEATURES = 40
+
+
+class EdgeClassifierV2(nn.Module):
+    def __init__(self, in_dim: int = NUM_FEATURES, hidden: int = 64):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden // 2),
+            nn.GELU(),
+            nn.Linear(hidden // 2, 1),
+        )
+
+    def forward(self, x):
+        return self.net(x).squeeze(-1)
+
+
+def _build_class_masks_and_corner_dt(good, views, dilate_px: int = 3):
+    """Per-view: dict img_id -> {edge_masks[7], corner_dt, H, W}."""
+    import cv2
+    from hoho2025.color_mappings import gestalt_color_mapping
+    out = {}
+    for gest_pil, depth_pil, img_id in zip(
+        good["gestalt"], good["depth"], good["image_ids"]
+    ):
+        if img_id not in views:
+            continue
+        depth_np = np.array(depth_pil)
+        H, W = depth_np.shape[:2]
+        gest_np = np.array(gest_pil.resize((W, H))).astype(np.uint8)
+
+        # Edge class masks: each cls -> bool mask
+        edge_masks = []
+        for cls in EDGE_CLASSES:
+            color = np.array(gestalt_color_mapping[cls])
+            m = cv2.inRange(gest_np, color - 0.5, color + 0.5)
+            if dilate_px > 0:
+                k = 2 * dilate_px + 1
+                m = cv2.dilate(m, np.ones((k, k), np.uint8), iterations=1)
+            edge_masks.append(m > 0)
+
+        # Corner DT: distance to nearest pixel of ANY point class
+        corner_mask = np.zeros((H, W), dtype=np.uint8)
+        for cls in POINT_CLASSES:
+            color = np.array(gestalt_color_mapping[cls])
+            corner_mask |= cv2.inRange(gest_np, color - 0.5, color + 0.5)
+        if corner_mask.sum() > 0:
+            corner_dt = cv2.distanceTransform(255 - corner_mask, cv2.DIST_L2, 5)
+            corner_dt = np.minimum(corner_dt, 30.0).astype(np.float32)
+        else:
+            corner_dt = np.full((H, W), 30.0, dtype=np.float32)
+
+        # Edge-union DT (for midpoint feature)
+        edge_union = np.zeros((H, W), dtype=np.uint8)
+        for cls in EDGE_CLASSES:
+            color = np.array(gestalt_color_mapping[cls])
+            edge_union |= cv2.inRange(gest_np, color - 0.5, color + 0.5)
+        if edge_union.sum() > 0:
+            edge_dt = cv2.distanceTransform(255 - edge_union, cv2.DIST_L2, 5)
+            edge_dt = np.minimum(edge_dt, 30.0).astype(np.float32)
+        else:
+            edge_dt = np.full((H, W), 30.0, dtype=np.float32)
+
+        out[img_id] = {
+            "edge_masks": edge_masks,  # list of 7 bool (H,W)
+            "corner_dt": corner_dt,
+            "edge_dt": edge_dt,
+            "H": H, "W": W,
+        }
+    return out
+
+
+def _edge_features_v2(pv, edges, sample, sample_steps: int = 24):
+    """Per-edge feature vectors (E, 40). All zeros on failure."""
+    from hoho2025.example_solutions import convert_entry_to_human_readable
+    from mvs_utils import collect_views, project_world_to_image
+    from scipy.spatial import cKDTree
+
+    E = len(edges)
+    feats = np.zeros((E, NUM_FEATURES), dtype=np.float32)
+    if E == 0:
+        return feats
+
+    pv_arr = np.asarray(pv, dtype=np.float64)
+
+    # Vertex degrees
+    deg = np.zeros(len(pv_arr), dtype=np.int32)
+    for a, b in edges:
+        if 0 <= a < len(pv_arr): deg[a] += 1
+        if 0 <= b < len(pv_arr): deg[b] += 1
+
+    # Edge geometry
+    lens = np.zeros(E, dtype=np.float32)
+    cos_z = np.zeros(E, dtype=np.float32)
+    for i, (a, b) in enumerate(edges):
+        if a >= len(pv_arr) or b >= len(pv_arr):
+            continue
+        d = pv_arr[b] - pv_arr[a]
+        n = float(np.linalg.norm(d))
+        lens[i] = n
+        if n > 1e-6:
+            cos_z[i] = float(d[2] / n)
+    median_len = float(np.median(lens[lens > 0])) if (lens > 0).any() else 1.0
+
+    # Geometric core (12 dims)
+    for i, (a, b) in enumerate(edges):
+        if a >= len(pv_arr) or b >= len(pv_arr):
+            continue
+        L = lens[i]
+        feats[i, 0] = float(np.log(L + 1e-3))
+        feats[i, 1] = cos_z[i]
+        feats[i, 2] = float(abs(cos_z[i]))
+        feats[i, 3] = float(deg[a])
+        feats[i, 4] = float(deg[b])
+        # 5-9 filled later (image features)
+        # 10-11 colmap (filled later)
+
+    try:
+        good = convert_entry_to_human_readable(sample)
+        colmap_rec = good.get("colmap") or good.get("colmap_binary")
+        if colmap_rec is None:
+            return feats
+        views = collect_views(colmap_rec, good["image_ids"])
+        if not views:
+            return feats
+        per_view = _build_class_masks_and_corner_dt(good, views)
+        if not per_view:
+            return feats
+
+        t_vec = np.linspace(0.0, 1.0, sample_steps)
+
+        # Colmap tree (for cdist)
+        pts = []
+        if hasattr(colmap_rec, 'points3D'):
+            for p in colmap_rec.points3D.values():
+                pts.append(p.xyz)
+        c_tree = cKDTree(np.asarray(pts, dtype=np.float32)) if pts else None
+
+        for i, (u, vv) in enumerate(edges):
+            u, vv = int(u), int(vv)
+            if u >= len(pv_arr) or vv >= len(pv_arr) or u == vv:
+                continue
+            endpoints = np.stack([pv_arr[u], pv_arr[vv]])
+
+            per_class_supp = []   # (n_views, 7)
+            any_edge_supp = []    # (n_views,)
+            proj_lens = []
+            corner_dist_a = []
+            corner_dist_b = []
+            close_a = 0
+            close_b = 0
+            mid_edge_dist = []
+            cam_z = []
+
+            for img_id, view in views.items():
+                if img_id not in per_view:
+                    continue
+                pv_view = per_view[img_id]
+                H, W = pv_view["H"], pv_view["W"]
+                uv, z = project_world_to_image(view["P"], endpoints)
+                if z[0] <= 0 or z[1] <= 0:
+                    continue
+                if not (0 <= uv[0,0] < W and 0 <= uv[0,1] < H
+                        and 0 <= uv[1,0] < W and 0 <= uv[1,1] < H):
+                    continue
+                cam_z.append(0.5 * (z[0] + z[1]))
+                proj_lens.append(float(np.linalg.norm(uv[1] - uv[0])))
+
+                xs = uv[0,0] + t_vec * (uv[1,0] - uv[0,0])
+                ys = uv[0,1] + t_vec * (uv[1,1] - uv[0,1])
+                xs_i = np.clip(xs.astype(np.int32), 0, W - 1)
+                ys_i = np.clip(ys.astype(np.int32), 0, H - 1)
+
+                # Per-class fractions
+                this_view_class = np.zeros(7, dtype=np.float32)
+                for ci in range(7):
+                    this_view_class[ci] = float(pv_view["edge_masks"][ci][ys_i, xs_i].sum()) / sample_steps
+                per_class_supp.append(this_view_class)
+                any_edge_supp.append(float(this_view_class.max()))  # max-class frac as "any edge"
+
+                # Endpoint corner distances
+                cd_a = float(pv_view["corner_dt"][int(uv[0,1]), int(uv[0,0])])
+                cd_b = float(pv_view["corner_dt"][int(uv[1,1]), int(uv[1,0])])
+                corner_dist_a.append(cd_a)
+                corner_dist_b.append(cd_b)
+                if cd_a < 10.0: close_a += 1
+                if cd_b < 10.0: close_b += 1
+
+                # Midpoint edge distance
+                mx, my = int(0.5 * (uv[0,0] + uv[1,0])), int(0.5 * (uv[0,1] + uv[1,1]))
+                mx = max(0, min(mx, W - 1))
+                my = max(0, min(my, H - 1))
+                mid_edge_dist.append(float(pv_view["edge_dt"][my, mx]))
+
+            if per_class_supp:
+                arr = np.asarray(per_class_supp)  # (n_views, 7)
+                feats[i, 12:19] = arr.mean(axis=0)
+                feats[i, 19:26] = arr.max(axis=0)
+                feats[i, 26] = float(np.mean(corner_dist_a))
+                feats[i, 27] = float(np.mean(corner_dist_b))
+                feats[i, 28] = float(close_a)
+                feats[i, 29] = float(close_b)
+                feats[i, 30] = float(np.mean(mid_edge_dist))
+                any_arr = np.asarray(any_edge_supp)
+                feats[i, 31] = float((any_arr > 0.3).sum())
+                feats[i, 32] = float((any_arr > 0.6).sum())
+                feats[i, 33] = float(np.log(max(np.mean(proj_lens), 1.0)))
+                feats[i, 34] = float(any_arr.std())
+                feats[i, 35] = float(np.mean(cam_z))
+
+            # Colmap distance
+            if c_tree is not None:
+                a3d, b3d = pv_arr[u], pv_arr[vv]
+                m3d = (a3d + b3d) * 0.5
+                d_a, _ = c_tree.query(a3d)
+                d_b, _ = c_tree.query(b3d)
+                d_m, _ = c_tree.query(m3d)
+                feats[i, 36] = float(min(d_a, d_b))
+                feats[i, 37] = float(d_m)
+
+            feats[i, 38] = float(lens[i] / max(median_len, 1e-3))
+            feats[i, 39] = float(len(pv_arr) / max(median_len, 1.0))  # graph density hint
+
+    except Exception:
+        pass
+
+    return feats
+
+
+def extract_features_v2(pv, pe, sample):
+    edges = [(int(a), int(b)) for a, b in pe]
+    return _edge_features_v2(pv, edges, sample)
+
+
+def label_edges_vs_gt(pv, pe, gt_v, gt_e, match_radius: float = 0.5):
+    pv_arr = np.asarray(pv, dtype=np.float32)
+    gt_v_arr = np.asarray(gt_v, dtype=np.float32)
+    labels = np.zeros(len(pe), dtype=np.float32)
+    if pv_arr.shape[0] < 2 or len(pe) == 0 or gt_v_arr.shape[0] < 2 or len(gt_e) == 0:
+        return labels
+    gt_pairs = []
+    for a, b in gt_e:
+        a, b = int(a), int(b)
+        if 0 <= a < len(gt_v_arr) and 0 <= b < len(gt_v_arr):
+            gt_pairs.append((gt_v_arr[a], gt_v_arr[b]))
+    if not gt_pairs:
+        return labels
+    r2 = match_radius * match_radius
+    for i, (u, vv) in enumerate(pe):
+        u, vv = int(u), int(vv)
+        if u >= len(pv_arr) or vv >= len(pv_arr):
+            continue
+        pa, pb = pv_arr[u], pv_arr[vv]
+        for ga, gb in gt_pairs:
+            d1 = ((pa - ga) ** 2).sum() + ((pb - gb) ** 2).sum()
+            d2 = ((pa - gb) ** 2).sum() + ((pb - ga) ** 2).sum()
+            if min(d1, d2) <= 2.0 * r2:
+                labels[i] = 1.0
+                break
+    return labels
+
+
+def classify_edges_v2(pv, pe, sample, classifier, threshold: float = 0.5,
+                      feature_mean=None, feature_std=None,
+                      min_keep_frac: float = 0.7):
+    try:
+        if len(pe) == 0:
+            return pv, pe
+        feats = extract_features_v2(pv, pe, sample)
+        if feature_mean is not None and feature_std is not None:
+            feats = (feats - feature_mean) / (feature_std + 1e-6)
+        with torch.no_grad():
+            x = torch.tensor(feats, dtype=torch.float32)
+            scores = torch.sigmoid(classifier(x)).numpy()
+
+        keep_mask = scores >= threshold
+        min_keep = max(1, int(np.ceil(min_keep_frac * len(pe))))
+        if keep_mask.sum() < min_keep:
+            top_idx = np.argsort(-scores)[:min_keep]
+            keep_mask = np.zeros_like(keep_mask)
+            keep_mask[top_idx] = True
+
+        keep_edges = [pe[i] for i in range(len(pe)) if keep_mask[i]]
+        if len(keep_edges) < 1:
+            return pv, pe
+
+        from edge_2d_filter import drop_orphan_vertices
+        return drop_orphan_vertices(np.asarray(pv), keep_edges)
+    except Exception:
+        return pv, pe
+
+
+def load_classifier_v2(path: str, device: str = "cpu"):
+    blob = torch.load(path, map_location=device, weights_only=False)
+    m = EdgeClassifierV2(in_dim=blob.get("in_dim", NUM_FEATURES),
+                         hidden=blob.get("hidden", 64))
+    m.load_state_dict(blob["model"])
+    m.to(device).eval()
+    return m, blob.get("feature_mean"), blob.get("feature_std")

edge_classifier_v4.py

@@ -0,0 +1,232 @@
+"""V4 edge classifier with DINOv2-pretrained patch features.
+
+DINOv2-small (22M params, frozen) encodes each gestalt view to a 16x16 grid
+of 384-dim semantic features. For each predicted edge, we bilinearly sample
+features at the projected edge midpoint in each view, then mean+max pool
+across views to get a 768-dim feature per edge. Concatenated with v2's
+40-dim geometric+mask features = 808-dim input to MLP.
+
+DINOv2 patch size is 14, input 224x224 → 16x16 patches. We feed gestalt RGB
+resized to 224x224.
+
+Falls back to no-op on any error.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from edge_classifier_v2 import (
+    NUM_FEATURES as V2_NUM, extract_features_v2,
+)
+
+DINO_FEAT_DIM = 384
+DINO_PATCH_SIZE = 14
+DINO_INPUT = 224
+DINO_GRID = DINO_INPUT // DINO_PATCH_SIZE  # 16
+EDGE_FEAT_DIM = DINO_FEAT_DIM * 2  # mean + max pool across views
+
+
+def get_dino_model(device="cpu"):
+    """Load DINOv2-small (frozen). Cached after first call."""
+    if not hasattr(get_dino_model, "_cache"):
+        get_dino_model._cache = {}
+    cache_key = str(device)
+    if cache_key not in get_dino_model._cache:
+        model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14', verbose=False)
+        model = model.to(device).eval()
+        for p in model.parameters():
+            p.requires_grad = False
+        get_dino_model._cache[cache_key] = model
+    return get_dino_model._cache[cache_key]
+
+
+class EdgeClassifierV4(nn.Module):
+    """Head: takes pre-pooled DINOv2 edge features + v2 geom features."""
+    def __init__(self, geom_dim: int = V2_NUM, edge_feat_dim: int = EDGE_FEAT_DIM,
+                 hidden: int = 128):
+        super().__init__()
+        in_dim = geom_dim + edge_feat_dim
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden),
+            nn.GELU(),
+            nn.Dropout(0.2),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Dropout(0.2),
+            nn.Linear(hidden, hidden // 2),
+            nn.GELU(),
+            nn.Linear(hidden // 2, 1),
+        )
+
+    def forward(self, geom_feats, edge_feats):
+        x = torch.cat([geom_feats, edge_feats], dim=1)
+        return self.net(x).squeeze(-1)
+
+
+@torch.no_grad()
+def _encode_views_with_dino(good, views, dino, device):
+    """Encode each view's gestalt image with DINOv2. Returns dict img_id -> (16,16,384)."""
+    out = {}
+    imgs = []
+    img_ids = []
+    Hs, Ws = {}, {}
+    for gest_pil, depth_pil, img_id in zip(
+        good["gestalt"], good["depth"], good["image_ids"]
+    ):
+        if img_id not in views:
+            continue
+        depth_np = np.array(depth_pil)
+        H, W = depth_np.shape[:2]
+        # Resize gestalt to 224x224 (DINOv2 input)
+        gest_resized = np.array(gest_pil.resize((DINO_INPUT, DINO_INPUT))).astype(np.float32) / 255.0
+        if gest_resized.ndim == 2:
+            gest_resized = np.stack([gest_resized]*3, axis=-1)
+        else:
+            gest_resized = gest_resized[..., :3]
+        # ImageNet normalization
+        gest_resized = (gest_resized - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225])
+        imgs.append(gest_resized.transpose(2, 0, 1))  # (3, 224, 224)
+        img_ids.append(img_id)
+        Hs[img_id] = H
+        Ws[img_id] = W
+        # Inject H,W for downstream code
+        views[img_id]["H"] = H
+        views[img_id]["W"] = W
+
+    if not imgs:
+        return out, Hs, Ws
+
+    batch = torch.tensor(np.stack(imgs), dtype=torch.float32, device=device)
+    # Forward (batch)
+    feats = dino.forward_features(batch)
+    patches = feats["x_norm_patchtokens"]  # (B, 256, 384)
+    patches = patches.reshape(len(imgs), DINO_GRID, DINO_GRID, DINO_FEAT_DIM).cpu().numpy()
+    for i, img_id in enumerate(img_ids):
+        out[img_id] = patches[i]  # (16, 16, 384)
+    return out, Hs, Ws
+
+
+def _bilinear_sample_grid(grid, u_norm, v_norm):
+    """Bilinearly sample (G, G, D) grid at (u, v) in [0, 1] coords. Returns (D,)."""
+    G = grid.shape[0]
+    u = u_norm * (G - 1)
+    v = v_norm * (G - 1)
+    x0 = int(np.floor(u)); x1 = min(x0 + 1, G - 1)
+    y0 = int(np.floor(v)); y1 = min(y0 + 1, G - 1)
+    x0 = max(0, x0); y0 = max(0, y0)
+    dx = u - x0; dy = v - y0
+    f00 = grid[y0, x0]
+    f01 = grid[y0, x1]
+    f10 = grid[y1, x0]
+    f11 = grid[y1, x1]
+    f0 = f00 * (1 - dx) + f01 * dx
+    f1 = f10 * (1 - dx) + f11 * dx
+    return f0 * (1 - dy) + f1 * dy
+
+
+def extract_features_v4(pv, pe, sample, dino, device="cpu"):
+    """Return (E, geom_dim), (E, edge_feat_dim) tensors."""
+    pv_arr = np.asarray(pv)
+    E = len(pe)
+    geom = extract_features_v2(pv, pe, sample)
+    edge_feats = np.zeros((E, EDGE_FEAT_DIM), dtype=np.float32)
+    if E == 0:
+        return geom, edge_feats
+
+    try:
+        from hoho2025.example_solutions import convert_entry_to_human_readable
+        from mvs_utils import collect_views, project_world_to_image
+        good = convert_entry_to_human_readable(sample)
+        colmap_rec = good.get("colmap") or good.get("colmap_binary")
+        if colmap_rec is None:
+            return geom, edge_feats
+        views = collect_views(colmap_rec, good["image_ids"])
+        if not views:
+            return geom, edge_feats
+        per_view, Hs, Ws = _encode_views_with_dino(good, views, dino, device)
+        if not per_view:
+            return geom, edge_feats
+
+        for i, (u, vv) in enumerate(pe):
+            u, vv = int(u), int(vv)
+            if u >= len(pv_arr) or vv >= len(pv_arr) or u == vv:
+                continue
+            endpoints = np.stack([pv_arr[u], pv_arr[vv]])
+            per_view_feats = []
+            for img_id, view in views.items():
+                if img_id not in per_view:
+                    continue
+                H, W = Hs[img_id], Ws[img_id]
+                uv, z = project_world_to_image(view["P"], endpoints)
+                if z[0] <= 0 or z[1] <= 0:
+                    continue
+                # Midpoint
+                mx, my = 0.5 * (uv[0, 0] + uv[1, 0]), 0.5 * (uv[0, 1] + uv[1, 1])
+                if not (0 <= mx < W and 0 <= my < H):
+                    continue
+                # Sample DINOv2 grid at (mx/W, my/H)
+                feat = _bilinear_sample_grid(per_view[img_id], mx / max(W - 1, 1), my / max(H - 1, 1))
+                per_view_feats.append(feat)
+            if per_view_feats:
+                arr = np.asarray(per_view_feats)  # (V, 384)
+                mean_f = arr.mean(axis=0)
+                max_f = arr.max(axis=0)
+                edge_feats[i] = np.concatenate([mean_f, max_f])
+    except Exception:
+        pass
+
+    return geom, edge_feats
+
+
+def label_edges_vs_gt(pv, pe, gt_v, gt_e, match_radius: float = 0.4):
+    from edge_classifier_v2 import label_edges_vs_gt as _f
+    return _f(pv, pe, gt_v, gt_e, match_radius)
+
+
+def classify_edges_v4(pv, pe, sample, classifier, dino, device="cpu",
+                     threshold: float = 0.5,
+                     feature_mean=None, feature_std=None,
+                     edge_feat_mean=None, edge_feat_std=None,
+                     min_keep_frac: float = 0.7):
+    try:
+        if len(pe) == 0:
+            return pv, pe
+        geom, edge_feats = extract_features_v4(pv, pe, sample, dino, device=device)
+        if feature_mean is not None and feature_std is not None:
+            geom = (geom - feature_mean) / (feature_std + 1e-6)
+        if edge_feat_mean is not None and edge_feat_std is not None:
+            edge_feats = (edge_feats - edge_feat_mean) / (edge_feat_std + 1e-6)
+        with torch.no_grad():
+            g = torch.tensor(geom, dtype=torch.float32, device=device)
+            e = torch.tensor(edge_feats, dtype=torch.float32, device=device)
+            scores = torch.sigmoid(classifier(g, e)).cpu().numpy()
+
+        keep_mask = scores >= threshold
+        min_keep = max(1, int(np.ceil(min_keep_frac * len(pe))))
+        if keep_mask.sum() < min_keep:
+            top_idx = np.argsort(-scores)[:min_keep]
+            keep_mask = np.zeros_like(keep_mask)
+            keep_mask[top_idx] = True
+
+        keep_edges = [pe[i] for i in range(len(pe)) if keep_mask[i]]
+        if len(keep_edges) < 1:
+            return pv, pe
+
+        from edge_2d_filter import drop_orphan_vertices
+        return drop_orphan_vertices(np.asarray(pv), keep_edges)
+    except Exception:
+        return pv, pe
+
+
+def load_classifier_v4(path: str, device: str = "cpu"):
+    blob = torch.load(path, map_location=device, weights_only=False)
+    m = EdgeClassifierV4(geom_dim=V2_NUM, edge_feat_dim=EDGE_FEAT_DIM,
+                         hidden=blob.get("hidden", 128))
+    m.load_state_dict(blob["model"])
+    m.to(device).eval()
+    return (m, blob.get("feature_mean"), blob.get("feature_std"),
+            blob.get("edge_feat_mean"), blob.get("edge_feat_std"))

local_eval.py

@@ -66,6 +66,54 @@ def parse_args():
                    help="hungarian TTA: drop anchor segments without this many supporting passes")
     p.add_argument("--tta-seeds", type=str, default="2718,31415,42",
                    help="comma-separated priority-sample seeds for TTA")
+    p.add_argument("--tracks-only", action="store_true",
+                   help="output ONLY the triangulation tracks (debug: baseline of tracks alone)")
+    p.add_argument("--fallback-to-tracks-when", type=str, default="",
+                   help="fallback to tracks-only when pred_v > X and track_v < Y, e.g. 20,8")
+    p.add_argument("--hallu-filter", action="store_true",
+                   help="filter vertices lacking BOTH COLMAP and gestalt-corner support")
+    p.add_argument("--hallu-colmap-radius", type=float, default=0.8,
+                   help="hallucination filter: COLMAP support radius (meters)")
+    p.add_argument("--hallu-gestalt-px", type=float, default=20.0,
+                   help="hallucination filter: gestalt corner pixel radius")
+    p.add_argument("--hallu-min-views", type=int, default=1,
+                   help="hallucination filter: min views with gestalt support")
+    p.add_argument("--ensemble", type=str, default="",
+                   help="comma-separated paths to ensemble checkpoints (additional to default checkpoint.pt)")
+    p.add_argument("--ensemble-min-passes", type=int, default=1,
+                   help="ensemble: min cross-pass agreement to keep anchor segments")
+    p.add_argument("--bundle-adjust", action="store_true",
+                   help="apply joint multi-view wireframe bundle adjustment")
+    p.add_argument("--ba-iter", type=int, default=50,
+                   help="bundle-adjust Adam iterations")
+    p.add_argument("--ba-lr", type=float, default=0.003,
+                   help="bundle-adjust learning rate")
+    p.add_argument("--ba-vertex-weight", type=float, default=1.0,
+                   help="bundle-adjust vertex (corner pixel) loss weight")
+    p.add_argument("--ba-edge-weight", type=float, default=0.5,
+                   help="bundle-adjust edge (edge pixel) loss weight")
+    p.add_argument("--ba-anchor-weight", type=float, default=200.0,
+                   help="bundle-adjust anchor regularization weight")
+    p.add_argument("--ba-max-move", type=float, default=0.4,
+                   help="bundle-adjust hard cap on vertex displacement (meters)")
+    p.add_argument("--tri-supplement", action="store_true",
+                   help="supplement sparse predictions with loose (min_views=2) tracks")
+    p.add_argument("--tri-sparse-threshold", type=int, default=5,
+                   help="tri-supplement: only activate when pred has < N vertices")
+    p.add_argument("--tri-merge-radius", type=float, default=0.7,
+                   help="tri-supplement: absorb loose vertex into pred if within radius")
+    p.add_argument("--edge-classifier", type=str, default="",
+                   help="path to edge_classifier.pt; if set, filter edges via learned model")
+    p.add_argument("--edge-classifier-v2", type=str, default="",
+                   help="path to edge_classifier_v2.pt; if set, use v2 with image-mask features")
+    p.add_argument("--edge-classifier-v3", type=str, default="",
+                   help="path to edge_classifier_v3.pt; v3 = CNN patches + v2 features")
+    p.add_argument("--edge-classifier-v4", type=str, default="",
+                   help="path to edge_classifier_v4.pt; v4 = DINOv2 features + v2 features")
+    p.add_argument("--edge-class-thresh", type=float, default=0.5,
+                   help="edge classifier: keep edges with P(keep) >= threshold")
+    p.add_argument("--edge-class-min-keep", type=float, default=0.5,
+                   help="edge classifier: never drop more than (1 - this) of edges")
     p.add_argument("--edge-fill", action="store_true",
                    help="enable edge filling from 2D mask evidence")
     p.add_argument("--fill-min-views", type=int, default=2,
@@ -81,7 +129,8 @@ def predict_one(sample, model, device, cfg, rng,
                 use_tracks=True, use_2d_filter=True, orphan_only=False,
                 strict_no_support=False, vertex_refine=False,
                 refine_kwargs=None, edge_fill=False, fill_kwargs=None,
-                tta=False, tta_seeds=None):
+                tta=False, tta_seeds=None,
+                tracks_only=False, fallback_tracks=None):
     """Run the full inference pipeline on one sample. Returns (pv, pe, diag)."""
     diag = {"colmap": -1, "fused": 0, "track_v": 0, "track_e": 0,
             "pred_v": 0, "pred_e": 0, "2dfilt_in": 0, "2dfilt_out": 0,
@@ -96,7 +145,24 @@ def predict_one(sample, model, device, cfg, rng,
     except Exception:
         pass
 
-    if tta:
+    if getattr(predict_one, "_ensemble_models", None):
+        try:
+            from ensemble import predict_sample_ensemble
+            # If --tta and --ensemble both set, use TTA seeds; else single seed.
+            if tta:
+                seeds = tta_seeds or (2718, 31415, 42)
+            else:
+                seeds = (2718,)
+            pred_v, pred_e = predict_sample_ensemble(
+                sample, cfg, predict_one._ensemble_models, device,
+                seeds=tuple(seeds),
+                min_passes_for_keep=getattr(predict_one, "_ensemble_min_passes", 1),
+            )
+            diag["fused"] = -1
+        except Exception as e:
+            diag["status"] = f"ensemble_failed:{type(e).__name__}"
+            return *script.empty_solution(), diag
+    elif tta:
         try:
             seeds = tta_seeds or (2718, 31415, 42)
             tta_method = (
@@ -142,6 +208,22 @@ def predict_one(sample, model, device, cfg, rng,
         except Exception as e:
             diag["status"] = f"track_failed:{type(e).__name__}"
 
+    # Sparse-scene supplement: when pred remains tiny after model + min_views=3
+    # tracks, fall back to min_views=2 loose tracks.
+    if getattr(predict_one, "_tri_supplement", False):
+        try:
+            from triangulate_supplement import supplement_sparse_with_loose_tracks
+            v_before = len(pred_v) if hasattr(pred_v, '__len__') else 0
+            pred_v, pred_e = supplement_sparse_with_loose_tracks(
+                pred_v, pred_e, sample,
+                sparse_threshold=getattr(predict_one, "_tri_sparse_threshold", 5),
+                merge_radius=getattr(predict_one, "_tri_merge_radius", 0.7),
+            )
+            v_after = len(pred_v) if hasattr(pred_v, '__len__') else 0
+            diag["tri_added"] = v_after - v_before
+        except Exception as e:
+            diag["status"] = f"tri_failed:{type(e).__name__}"
+
     diag["2dfilt_in"] = len(pred_e) if hasattr(pred_e, '__len__') else 0
     # Vertex refinement runs FIRST (refines vertex positions while orphan is still present;
     # orphan/2d-filter then cleans up afterwards).
@@ -184,6 +266,40 @@ def predict_one(sample, model, device, cfg, rng,
             diag["status"] = f"2dfilt_failed:{type(e).__name__}"
     diag["2dfilt_out"] = len(pred_e) if hasattr(pred_e, '__len__') else 0
 
+    # Edge classifier: learned keep/drop on top of post-filter edges.
+    ec = getattr(predict_one, "_edge_classifier", None)
+    if ec is not None:
+        try:
+            e_before = len(pred_e) if hasattr(pred_e, '__len__') else 0
+            ver = ec.get("version", 1)
+            if ver == 4:
+                from edge_classifier_v4 import classify_edges_v4
+                pred_v, pred_e = classify_edges_v4(
+                    pred_v, pred_e, sample,
+                    ec["model"], ec["dino"], device=ec["dino_device"],
+                    threshold=ec["threshold"],
+                    feature_mean=ec["mean"], feature_std=ec["std"],
+                    edge_feat_mean=ec["edge_feat_mean"], edge_feat_std=ec["edge_feat_std"],
+                    min_keep_frac=ec["min_keep_frac"],
+                )
+            else:
+                if ver == 3:
+                    from edge_classifier_v3 import classify_edges_v3 as _cls_fn
+                elif ver == 2:
+                    from edge_classifier_v2 import classify_edges_v2 as _cls_fn
+                else:
+                    from edge_classifier import classify_edges as _cls_fn
+                pred_v, pred_e = _cls_fn(
+                    pred_v, pred_e, sample,
+                    ec["model"], threshold=ec["threshold"],
+                    feature_mean=ec["mean"], feature_std=ec["std"],
+                    min_keep_frac=ec["min_keep_frac"],
+                )
+            diag["ec_kept"] = (len(pred_e) if hasattr(pred_e, '__len__') else 0)
+            diag["ec_dropped"] = e_before - diag["ec_kept"]
+        except Exception as e:
+            diag["status"] = f"ec_failed:{type(e).__name__}"
+
     if edge_fill:
         e_before = len(pred_e) if hasattr(pred_e, '__len__') else 0
         try:
@@ -195,6 +311,57 @@ def predict_one(sample, model, device, cfg, rng,
         except Exception as e:
             diag["status"] = f"fill_failed:{type(e).__name__}"
 
+    if getattr(predict_one, "_bundle_adjust", False):
+        try:
+            from bundle_wireframe import bundle_adjust_wireframe
+            ba_kwargs = getattr(predict_one, "_ba_kwargs", {})
+            pred_v, pred_e = bundle_adjust_wireframe(
+                pred_v, pred_e, sample, **ba_kwargs)
+        except Exception as e:
+            diag["status"] = f"ba_failed:{type(e).__name__}"
+
+    # Optional: hallucination filter (drop vertices lacking both COLMAP and gestalt support)
+    if getattr(predict_one, "_hallu_filter", False):
+        try:
+            from hallucination_filter import filter_hallucinated_vertices
+            n_before = len(pred_v) if hasattr(pred_v, '__len__') else 0
+            pred_v, pred_e = filter_hallucinated_vertices(
+                pred_v, pred_e, sample,
+                colmap_radius=getattr(predict_one, "_hallu_colmap_radius", 0.8),
+                gestalt_radius_px=getattr(predict_one, "_hallu_gestalt_px", 20.0),
+                min_views_with_gestalt=getattr(predict_one, "_hallu_min_views", 1),
+            )
+            n_after = len(pred_v) if hasattr(pred_v, '__len__') else 0
+            diag["hallu_dropped"] = n_before - n_after
+        except Exception as e:
+            diag["status"] = f"hallu_failed:{type(e).__name__}"
+
+    # Optional: replace prediction with tracks-only on hard scenes
+    if tracks_only:
+        # Use the tracks computed earlier (regardless of pred quality)
+        try:
+            from triangulation import predict_wireframe_tracks
+            track_v, track_e = predict_wireframe_tracks(sample, min_views=3)
+            if track_v is not None and track_e is not None and len(track_v) >= 2 and len(track_e) >= 1:
+                pred_v = np.asarray(track_v, dtype=np.float32)
+                pred_e = list(track_e)
+                diag["status"] = "tracks_only_forced"
+        except Exception:
+            pass
+    elif fallback_tracks is not None:
+        pred_v_thresh, track_v_thresh = fallback_tracks  # tuple
+        n_pv = len(pred_v) if hasattr(pred_v, '__len__') else 0
+        if n_pv > pred_v_thresh and diag.get("track_v", 0) < track_v_thresh:
+            try:
+                from triangulation import predict_wireframe_tracks
+                track_v, track_e = predict_wireframe_tracks(sample, min_views=3)
+                if track_v is not None and track_e is not None and len(track_v) >= 2 and len(track_e) >= 1:
+                    pred_v = np.asarray(track_v, dtype=np.float32)
+                    pred_e = list(track_e)
+                    diag["status"] = "fallback_to_tracks"
+            except Exception:
+                pass
+
     diag["pred_v"] = len(pred_v) if hasattr(pred_v, '__len__') else 0
     diag["pred_e"] = len(pred_e) if hasattr(pred_e, '__len__') else 0
     return pred_v, pred_e, diag
@@ -222,6 +389,72 @@ def main():
     model = script.load_model(ckpt_path, device)
     print(f"Model: {sum(p.numel() for p in model.parameters()):,} params")
 
+    ensemble_models = None
+    if args.ensemble:
+        from ensemble import load_two_checkpoints
+        extra_paths = [p.strip() for p in args.ensemble.split(",") if p.strip()]
+        ensemble_models = [model] + load_two_checkpoints(extra_paths, device)
+        print(f"Ensemble: {len(ensemble_models)} models")
+
+    if args.edge_classifier:
+        from edge_classifier import load_classifier
+        ec_model, ec_mean, ec_std = load_classifier(args.edge_classifier, device="cpu")
+        print(f"Edge classifier loaded from {args.edge_classifier} "
+              f"(thresh={args.edge_class_thresh}, min_keep={args.edge_class_min_keep})")
+        predict_one._edge_classifier_loaded = {
+            "model": ec_model,
+            "mean": ec_mean.cpu().numpy() if hasattr(ec_mean, 'cpu') else ec_mean,
+            "std": ec_std.cpu().numpy() if hasattr(ec_std, 'cpu') else ec_std,
+            "threshold": args.edge_class_thresh,
+            "min_keep_frac": args.edge_class_min_keep,
+            "version": 1,
+        }
+
+    if args.edge_classifier_v2:
+        from edge_classifier_v2 import load_classifier_v2
+        ec_model, ec_mean, ec_std = load_classifier_v2(args.edge_classifier_v2, device="cpu")
+        print(f"Edge classifier V2 loaded from {args.edge_classifier_v2}")
+        predict_one._edge_classifier_loaded = {
+            "model": ec_model,
+            "mean": ec_mean.cpu().numpy() if hasattr(ec_mean, 'cpu') else ec_mean,
+            "std": ec_std.cpu().numpy() if hasattr(ec_std, 'cpu') else ec_std,
+            "threshold": args.edge_class_thresh,
+            "min_keep_frac": args.edge_class_min_keep,
+            "version": 2,
+        }
+
+    if args.edge_classifier_v3:
+        from edge_classifier_v3 import load_classifier_v3
+        ec_model, ec_mean, ec_std = load_classifier_v3(args.edge_classifier_v3, device="cpu")
+        print(f"Edge classifier V3 (CNN) loaded from {args.edge_classifier_v3}")
+        predict_one._edge_classifier_loaded = {
+            "model": ec_model,
+            "mean": ec_mean.cpu().numpy() if hasattr(ec_mean, 'cpu') else ec_mean,
+            "std": ec_std.cpu().numpy() if hasattr(ec_std, 'cpu') else ec_std,
+            "threshold": args.edge_class_thresh,
+            "min_keep_frac": args.edge_class_min_keep,
+            "version": 3,
+        }
+
+    if args.edge_classifier_v4:
+        from edge_classifier_v4 import load_classifier_v4, get_dino_model
+        ec_model, g_mean, g_std, e_mean, e_std = load_classifier_v4(args.edge_classifier_v4, device="cpu")
+        # DINO runs on the inference device for speed
+        dino = get_dino_model(device=device)
+        print(f"Edge classifier V4 (DINOv2) loaded from {args.edge_classifier_v4}")
+        predict_one._edge_classifier_loaded = {
+            "model": ec_model,
+            "dino": dino,
+            "dino_device": device,
+            "mean": g_mean.cpu().numpy() if hasattr(g_mean, 'cpu') else g_mean,
+            "std": g_std.cpu().numpy() if hasattr(g_std, 'cpu') else g_std,
+            "edge_feat_mean": e_mean.cpu().numpy() if hasattr(e_mean, 'cpu') else e_mean,
+            "edge_feat_std": e_std.cpu().numpy() if hasattr(e_std, 'cpu') else e_std,
+            "threshold": args.edge_class_thresh,
+            "min_keep_frac": args.edge_class_min_keep,
+            "version": 4,
+        }
+
     if args.conf_thresh is not None:
         print(f"Overriding script.CONF_THRESH: {script.CONF_THRESH} -> {args.conf_thresh}")
         script.CONF_THRESH = args.conf_thresh
@@ -271,6 +504,29 @@ def main():
             tta_seeds_tuple = tuple(int(s) for s in args.tta_seeds.split(","))
             predict_one._tta_hungarian = args.tta_hungarian
             predict_one._tta_min_passes = args.tta_min_passes
+            predict_one._hallu_filter = args.hallu_filter
+            predict_one._hallu_colmap_radius = args.hallu_colmap_radius
+            predict_one._hallu_gestalt_px = args.hallu_gestalt_px
+            predict_one._hallu_min_views = args.hallu_min_views
+            predict_one._ensemble_models = ensemble_models
+            predict_one._ensemble_min_passes = args.ensemble_min_passes
+            predict_one._bundle_adjust = args.bundle_adjust
+            predict_one._ba_kwargs = {
+                "n_iter": args.ba_iter,
+                "lr": args.ba_lr,
+                "vertex_weight": args.ba_vertex_weight,
+                "edge_weight": args.ba_edge_weight,
+                "anchor_weight": args.ba_anchor_weight,
+                "max_move_meters": args.ba_max_move,
+            }
+            predict_one._tri_supplement = args.tri_supplement
+            predict_one._tri_sparse_threshold = args.tri_sparse_threshold
+            predict_one._tri_merge_radius = args.tri_merge_radius
+            predict_one._edge_classifier = getattr(predict_one, "_edge_classifier_loaded", None)
+            fallback_tracks = None
+            if args.fallback_to_tracks_when:
+                pv_thr, tv_thr = args.fallback_to_tracks_when.split(",")
+                fallback_tracks = (int(pv_thr), int(tv_thr))
             pred_v, pred_e, diag = predict_one(
                 sample, model, device, cfg, rng,
                 use_tracks=not args.no_tracks,
@@ -282,7 +538,9 @@ def main():
                 edge_fill=args.edge_fill,
                 fill_kwargs=fill_kwargs,
                 tta=args.tta,
-                tta_seeds=tta_seeds_tuple)
+                tta_seeds=tta_seeds_tuple,
+                tracks_only=args.tracks_only,
+                fallback_tracks=fallback_tracks)
             if torch.backends.mps.is_available():
                 torch.mps.empty_cache()
 

script.py

@@ -391,6 +391,37 @@ if __name__ == "__main__":
     model = load_model(checkpoint_path, device)
     print(f"Model loaded: {sum(p.numel() for p in model.parameters()):,} params")
 
+    # Edge classifier (DINOv2 patch features + geometric features → P(keep edge)).
+    # Trained on 400 samples with match_radius=0.4. Best val acc 82.1%.
+    # Operating point: thresh=0.15, min_keep=0.85 — drop only the bottom ~15%
+    # of edges that the classifier is most confident are wrong. 200-sample local
+    # A/B: +0.0030 hss_mean (t=+1.26, 100 wins / 77 losses).
+    edge_classifier_bundle = None
+    edge_classifier_path = SCRIPT_DIR / "edge_classifier_v4_400.pt"
+    if edge_classifier_path.exists():
+        try:
+            from edge_classifier_v4 import load_classifier_v4, get_dino_model
+            ec_model, g_mean, g_std, e_mean, e_std = load_classifier_v4(
+                str(edge_classifier_path), device="cpu")
+            dino = get_dino_model(device=device)
+            edge_classifier_bundle = {
+                "model": ec_model,
+                "dino": dino,
+                "dino_device": device,
+                "mean": g_mean.cpu().numpy() if hasattr(g_mean, "cpu") else g_mean,
+                "std": g_std.cpu().numpy() if hasattr(g_std, "cpu") else g_std,
+                "edge_feat_mean": e_mean.cpu().numpy() if hasattr(e_mean, "cpu") else e_mean,
+                "edge_feat_std": e_std.cpu().numpy() if hasattr(e_std, "cpu") else e_std,
+                "threshold": 0.15,
+                "min_keep_frac": 0.85,
+            }
+            print(f"Edge classifier v4 loaded ({sum(p.numel() for p in ec_model.parameters()):,} head params + frozen DINOv2)")
+        except Exception as ec_err:
+            print(f"Edge classifier load failed: {ec_err}; running without")
+            edge_classifier_bundle = None
+    else:
+        print(f"No edge_classifier_v4_400.pt at {edge_classifier_path}; running without")
+
     # Optional: load 2nd checkpoint for ensemble inference
     ensemble_models = None
     if USE_ENSEMBLE:
@@ -550,6 +581,28 @@ if __name__ == "__main__":
                         print(f"  orphan drop failed for {order_id}: {filt_err}")
                     edges_after_2d = len(pred_e) if hasattr(pred_e, '__len__') else 0
 
+                    # Edge classifier v4 (DINOv2 + geom features): drop the bottom
+                    # ~15% of edges whose learned P(keep) is lowest. 200-sample
+                    # local A/B: +0.0030 hss_mean (t=+1.26).
+                    if edge_classifier_bundle is not None:
+                        try:
+                            from edge_classifier_v4 import classify_edges_v4
+                            from edge_2d_filter import drop_orphan_vertices as _drop_orph
+                            ec = edge_classifier_bundle
+                            pred_v, pred_e = classify_edges_v4(
+                                pred_v, pred_e, sample,
+                                ec["model"], ec["dino"], device=ec["dino_device"],
+                                threshold=ec["threshold"],
+                                feature_mean=ec["mean"], feature_std=ec["std"],
+                                edge_feat_mean=ec["edge_feat_mean"],
+                                edge_feat_std=ec["edge_feat_std"],
+                                min_keep_frac=ec["min_keep_frac"],
+                            )
+                            # Re-run orphan drop in case the classifier left orphans
+                            pred_v, pred_e = _drop_orph(pred_v, pred_e)
+                        except Exception as ec_err:
+                            print(f"  edge classifier v4 failed for {order_id}: {ec_err}")
+
                 except Exception as e:
                     import traceback
                     print(f"  Predict failed for {order_id}:\n{traceback.format_exc()}")