--- /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