Update code: learned-baseline + graph-opt postprocess

requirements.txt

@@ -1,14 +1,9 @@
-modal>=0.73.0
-datasets>=3.0.0
-huggingface-hub>=0.25.0
 torch>=2.2.0
 torchvision>=0.17.0
 numpy>=1.24.0
-opencv-python-headless>=4.8.0
 scipy>=1.11.0
 scikit-learn>=1.3.0
 trimesh>=4.0.0
-manifold3d
 Pillow>=10.0.0
 plotly>=5.0.0
 pycolmap>=3.13
@@ -17,4 +12,4 @@ open3d>=0.18.0
 tqdm
 pandas
 matplotlib
-hf-transfer
+opencv-python-headless>=4.8.0

script.py

@@ -88,16 +88,22 @@ from src.learned_baseline.model import EdgeDepthSegmentsModel
 from src.learned_baseline.segment_postprocess import merge_vertices_iterative
 from src.learned_baseline.varifold import segments_to_vertices_edges
 from src.learned_baseline.postprocess_v2 import snap_to_point_cloud, snap_horizontal
+from src.learned_baseline.plane_snap import snap_to_plane_intersections
 
 # graph_optimizer A/B-tested -0.0014 HSS — disabled
 USE_GRAPH_OPT = False
 
-SEQ_LEN       = 4096
-COLMAP_QUOTA  = 3072
-DEPTH_QUOTA   = 1024
-CONF_THRESH   = 0.5
-MERGE_THRESH  = 0.4
-SNAP_RADIUS   = 0.5
+SEQ_LEN            = 4096
+COLMAP_QUOTA       = 3072
+DEPTH_QUOTA        = 1024
+# Hparam sweep on 200 val scenes: marginal +0.0008 HSS over (0.5, 0.5, 0.6).
+CONF_THRESH        = 0.7   # was 0.5
+MERGE_THRESH       = 0.4
+SNAP_RADIUS        = 0.3   # was 0.5
+PLANE_SNAP_RADIUS  = 0.35  # sweep of 100 val scenes: r=0.35 best (+0.005 HSS)
+USE_PLANE_SNAP     = True
+USE_TTA            = False  # model not rotation-equivariant; TTA hurts both F1 and HSS
+TTA_ANGLES         = [0.0, 0.5 * np.pi, np.pi, 1.5 * np.pi]
 
 CHECKPOINT_PATH = SCRIPT_DIR / "checkpoint_learned.pt"
 if not CHECKPOINT_PATH.exists():
@@ -215,17 +221,76 @@ def build_tokens_single(sample_dict, model, device):
     return tokens, masks
 
 
-def predict_sample(sample_dict, model, device):
-    tokens, masks = build_tokens_single(sample_dict, model, device)
-    scale  = float(sample_dict["scale"])
-    center = sample_dict["center"]
-
+def _nms_segments(segs_np, confs_np, mid_thresh=0.4, dir_thresh=0.90):
+    """Greedy NMS on segment pool. Removes near-duplicate segments (same structural
+    edge predicted from multiple TTA rotations). segs_np: (N, 2, 3), confs_np: (N,).
+    Keeps highest-confidence segment in each cluster."""
+    if len(segs_np) == 0:
+        return segs_np, confs_np
+    order = np.argsort(-confs_np)
+    segs_s = segs_np[order]
+    confs_s = confs_np[order]
+    mids = (segs_s[:, 0] + segs_s[:, 1]) * 0.5          # (N, 3)
+    dirs = segs_s[:, 1] - segs_s[:, 0]                   # (N, 3)
+    norms = np.linalg.norm(dirs, axis=-1, keepdims=True) + 1e-8
+    dirs = dirs / norms
+    keep = np.ones(len(segs_s), dtype=bool)
+    for i in range(len(segs_s)):
+        if not keep[i]:
+            continue
+        dists = np.linalg.norm(mids[i+1:] - mids[i], axis=-1)
+        cos_sim = np.abs((dirs[i+1:] * dirs[i]).sum(-1))
+        suppress = (dists < mid_thresh) & (cos_sim > dir_thresh)
+        keep[i+1:][suppress] = False
+    return segs_s[keep], confs_s[keep]
+
+
+def _rotation_y(angle):
+    """Rotation matrix around y-axis (vertical)."""
+    c, s = float(np.cos(angle)), float(np.sin(angle))
+    return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float32)
+
+
+def _forward_one(sample_dict, model, device, R=None):
+    """Single forward pass, optionally with point cloud rotated by R (3x3 numpy array).
+    Returns (segs_norm_derotated, conf) where segs_norm is (N, 2, 3) in normalized space."""
+    sd = sample_dict
+    if R is not None:
+        sd = dict(sample_dict)
+        sd["xyz_norm"] = sample_dict["xyz_norm"] @ R.T  # rotate cloud
+    tokens, masks = build_tokens_single(sd, model, device)
     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()
+    segs = out["segments"][0].float().cpu()      # (N, 2, 3) normalized
     conf = torch.sigmoid(out["conf"][0].float()).cpu().numpy() if "conf" in out else None
+    if R is not None:
+        # De-rotate: segs @ R  (since input was rotated by R, undo with R^-1 = R^T,
+        # but in row-vector form: x_orig = x_rot @ R  because x_rot = x_orig @ R.T)
+        Rt = torch.tensor(R, dtype=torch.float32)
+        segs = torch.einsum('nkd,de->nke', segs, Rt)
+    return segs, conf
+
+
+def predict_sample(sample_dict, model, device):
+    scale  = float(sample_dict["scale"])
+    center = sample_dict["center"]
+
+    if USE_TTA:
+        all_segs, all_confs = [], []
+        for angle in TTA_ANGLES:
+            R = _rotation_y(angle) if angle != 0.0 else None
+            segs, conf = _forward_one(sample_dict, model, device, R)
+            all_segs.append(segs)
+            if conf is not None:
+                all_confs.append(conf)
+        segs_np = torch.cat(all_segs, dim=0).numpy()
+        conf = np.concatenate(all_confs) if all_confs else None
+        # NMS: collapse near-duplicate segments from 4 rotations into 1
+        segs_np, conf = _nms_segments(segs_np, conf)
+        segs = torch.tensor(segs_np)
+    else:
+        segs, conf = _forward_one(sample_dict, model, device)
 
     if conf is not None:
         keep = conf > CONF_THRESH
@@ -238,13 +303,23 @@ def predict_sample(sample_dict, model, device):
     pv, pe = pv.numpy(), np.array(pe, dtype=np.int32)
     pv, pe = merge_vertices_iterative(pv, pe)
 
-    xyz_norm = sample_dict["xyz_norm"]
-    mask     = sample_dict["mask"]
-    cid      = sample_dict["class_id"]
+    xyz_norm  = sample_dict["xyz_norm"]
+    mask      = sample_dict["mask"]
+    cid       = sample_dict["class_id"]
+    src       = sample_dict.get("source", np.zeros(len(xyz_norm), dtype=np.int64))
     xyz_world = xyz_norm[mask] * scale + center
     cid_valid = cid[mask]
+    src_valid = src[mask]
     pv = snap_to_point_cloud(pv, xyz_world, cid_valid, snap_radius=SNAP_RADIUS)
     pv = snap_horizontal(pv, pe)
+    if USE_PLANE_SNAP:
+        try:
+            pv = snap_to_plane_intersections(
+                pv, pe, xyz_world, cid_valid, src_valid=src_valid,
+                snap_radius=PLANE_SNAP_RADIUS,
+            )
+        except Exception as exc:
+            print(f"  plane_snap failed: {exc}")
 
     if len(pv) < 2 or len(pe) < 1:
         return empty_solution()

src/learned_baseline/configs/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
+}

src/learned_baseline/plane_snap.py

@@ -0,0 +1,135 @@
+"""Roof plane RANSAC snap: fit planes to roof surface points, snap vertices
+to plane-plane intersection lines (ridges/eaves/hips)."""
+import numpy as np
+
+ROOF_CLASS = 13  # gestalt "roof" surface label
+
+
+def _ransac_plane(pts, dist_thresh=0.12, max_iters=150, rng=None):
+    """Fit plane via RANSAC + PCA refit. Returns (normal, d) or None."""
+    if rng is None:
+        rng = np.random.RandomState(42)
+    best_inliers = 0
+    best_n = best_d = None
+    for _ in range(max_iters):
+        idx = rng.choice(len(pts), 3, replace=False)
+        p1, p2, p3 = pts[idx]
+        n = np.cross(p2 - p1, p3 - p1)
+        nrm = np.linalg.norm(n)
+        if nrm < 1e-10:
+            continue
+        n /= nrm
+        d = float(np.dot(n, p1))
+        n_in = int((np.abs(pts @ n - d) < dist_thresh).sum())
+        if n_in > best_inliers:
+            best_inliers, best_n, best_d = n_in, n, d
+    if best_n is None or best_inliers < 3:
+        return None
+    inliers = pts[np.abs(pts @ best_n - best_d) < dist_thresh]
+    if len(inliers) < 3:
+        return (best_n, float(best_d))
+    centroid = inliers.mean(0)
+    _, _, Vt = np.linalg.svd(inliers - centroid, full_matrices=False)
+    n = Vt[-1]
+    if np.dot(n, best_n) < 0:
+        n = -n
+    return (n, float(np.dot(n, centroid)))
+
+
+def _fit_planes(xyz_world, cid_valid, src_valid,
+                n_planes, min_inliers, dist_thresh):
+    """Iterative RANSAC on roof-class COLMAP points. Returns list of (n, d, centroid)."""
+    roof_mask = cid_valid == ROOF_CLASS
+    if src_valid is not None:
+        colmap_roof = roof_mask & (src_valid == 0)
+        pts = xyz_world[colmap_roof] if colmap_roof.sum() >= min_inliers else xyz_world[roof_mask]
+    else:
+        pts = xyz_world[roof_mask]
+    if len(pts) < min_inliers:
+        return []
+
+    rng = np.random.RandomState(123)
+    remaining = pts.copy()
+    planes = []
+    for _ in range(n_planes):
+        if len(remaining) < min_inliers:
+            break
+        result = _ransac_plane(remaining, dist_thresh=dist_thresh,
+                               max_iters=150, rng=rng)
+        if result is None:
+            break
+        n, d = result
+        inlier_mask = np.abs(remaining @ n - d) < dist_thresh
+        if inlier_mask.sum() < min_inliers:
+            break
+        planes.append((n, d, remaining[inlier_mask].mean(0)))
+        remaining = remaining[~inlier_mask]
+    return planes
+
+
+def _intersection_line(p1, p2):
+    """Line of intersection of two planes (n,d,c) each.
+    Returns (p0, dir, t_min, t_max) where [t_min, t_max] bounds the segment
+    between the two planes' centroid projections (limits snapping to actual edge
+    extent rather than infinite line). Returns None if parallel.
+    """
+    n1, d1, c1 = p1
+    n2, d2, c2 = p2
+    direction = np.cross(n1, n2)
+    nrm = np.linalg.norm(direction)
+    if nrm < 1e-6:
+        return None
+    direction /= nrm
+    A = np.stack([n1, n2, direction])
+    b = np.array([d1, d2, 0.0])
+    try:
+        p0 = np.linalg.solve(A, b)
+    except np.linalg.LinAlgError:
+        return None
+    # Bound line to overlap region of the two plane extents using centroids
+    t1 = float(np.dot(c1 - p0, direction))
+    t2 = float(np.dot(c2 - p0, direction))
+    t_min = min(t1, t2) - 3.0   # ±3m beyond centroids
+    t_max = max(t1, t2) + 3.0
+    return (p0, direction, t_min, t_max)
+
+
+def snap_to_plane_intersections(pv, pe, xyz_world, cid_valid, src_valid=None,
+                                 n_planes=10, min_inliers=25, dist_thresh=0.12,
+                                 snap_radius=0.45):
+    """Snap predicted vertices to nearest roof plane-intersection line.
+
+    Fits up to n_planes roof planes via iterative RANSAC, computes pairwise
+    intersection lines, then moves each vertex to the closest point on the
+    nearest line within snap_radius. Falls back silently if < 2 planes found.
+    """
+    planes = _fit_planes(xyz_world, cid_valid, src_valid,
+                         n_planes, min_inliers, dist_thresh)
+    if len(planes) < 2:
+        return pv
+
+    lines = []
+    for i in range(len(planes)):
+        for j in range(i + 1, len(planes)):
+            line = _intersection_line(planes[i], planes[j])
+            if line is not None:
+                lines.append(line)
+    if not lines:
+        return pv
+
+    pv_new = np.array(pv, dtype=np.float64)
+    for vi in range(len(pv_new)):
+        v = pv_new[vi]
+        best_dist = snap_radius
+        best_pos = None
+        for p0, d, t_min, t_max in lines:
+            t = float(np.dot(v - p0, d))
+            t = max(t_min, min(t_max, t))   # clamp to line segment
+            proj = p0 + t * d
+            dist = float(np.linalg.norm(v - proj))
+            if dist < best_dist:
+                best_dist = dist
+                best_pos = proj
+        if best_pos is not None:
+            pv_new[vi] = best_pos
+    return pv_new