src/learned_baseline/plane_snap.py

"""Structural plane fitting and vertex snapping for S23DR wireframes.

Class IDs (from cache_scenes.py STRUCTURAL_CLASSES, 0-indexed):
  0  apex            7  valley
  1  eave_end_point  8  flashing
  2  flashing_end_pt 9  step_flashing
  3  rake            10 roof          ← roof face surface
  4  ridge           11 other_house   ← walls, siding, doors, windows
  5  eave            12 non_house     ← unlabeled / ADE
  6  hip
"""
import numpy as np

ROOF_CLASS  = 10   # compressed class_id for "roof" face
WALL_CLASS  = 11   # compressed class_id for "other_house" (walls dominate)


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_iterative(pts, n_planes, min_inliers, dist_thresh):
    """Iterative RANSAC on pts. Returns list of (normal, d, centroid)."""
    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. Returns (p0, dir, t_min, t_max) or None."""
    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
    t1 = float(np.dot(c1 - p0, direction))
    t2 = float(np.dot(c2 - p0, direction))
    t_min = min(t1, t2) - 3.0
    t_max = max(t1, t2) + 3.0
    return (p0, direction, t_min, t_max)


def _plane_triple_intersection(p1, p2, p3):
    """Point where 3 planes (n,d,c) each intersect. Returns point or None."""
    n1, d1, _ = p1
    n2, d2, _ = p2
    n3, d3, _ = p3
    A = np.stack([n1, n2, n3])
    b = np.array([d1, d2, d3], dtype=np.float64)
    try:
        if abs(np.linalg.det(A)) < 1e-5:
            return None
        return np.linalg.solve(A, b)
    except np.linalg.LinAlgError:
        return None


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.

    Uses roof-class COLMAP points (class_id==10). Falls back if <2 planes found.
    """
    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 pv

    planes = _fit_planes_iterative(pts, 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))
            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


def snap_to_structural_intersections(pv, pe, xyz_world, cid_valid, src_valid=None,
                                      n_planes_wall=6, n_planes_roof=6,
                                      min_inliers=20, dist_thresh=0.12,
                                      snap_radius=0.5):
    """Snap vertices to structural plane triple-intersections.

    Fits planes to wall (class 11) and roof (class 10) COLMAP points, computes
    all 3-plane intersections, snaps each predicted vertex to the nearest one
    within snap_radius. Falls back silently if insufficient points/planes.

    3-plane intersections correspond to true wireframe vertices:
      wall-wall-floor → base corner
      wall-wall-ceiling → top corner
      wall-roof-roof → eave corner
      roof-roof-roof → ridge peak
    """
    pv = np.asarray(pv, dtype=np.float64)
    if len(pv) < 2:
        return pv

    # COLMAP points only (more reliable geometry)
    if src_valid is not None:
        colmap = src_valid == 0
        pts_all = xyz_world[colmap]
        cids_all = cid_valid[colmap]
    else:
        pts_all = xyz_world
        cids_all = cid_valid

    all_planes = []

    # Wall planes
    wall_pts = pts_all[cids_all == WALL_CLASS]
    if len(wall_pts) >= min_inliers:
        all_planes.extend(
            _fit_planes_iterative(wall_pts, n_planes_wall, min_inliers, dist_thresh)
        )

    # Roof planes
    roof_pts = pts_all[cids_all == ROOF_CLASS]
    if len(roof_pts) >= min_inliers:
        all_planes.extend(
            _fit_planes_iterative(roof_pts, n_planes_roof, min_inliers, dist_thresh)
        )

    if len(all_planes) < 3:
        return pv

    # Scene bounding box for filtering implausible intersections
    bb_min = pts_all.min(0) - 2.0
    bb_max = pts_all.max(0) + 2.0

    # All 3-plane intersections
    candidates = []
    n = len(all_planes)
    for i in range(n):
        for j in range(i + 1, n):
            for k in range(j + 1, n):
                pt = _plane_triple_intersection(all_planes[i], all_planes[j], all_planes[k])
                if pt is None:
                    continue
                if np.all(pt >= bb_min) and np.all(pt <= bb_max):
                    candidates.append(pt)

    if not candidates:
        return pv

    cand_arr = np.array(candidates)  # (M, 3)

    # Snap each vertex to nearest candidate within radius
    pv_new = pv.copy()
    for vi in range(len(pv_new)):
        dists = np.linalg.norm(cand_arr - pv_new[vi], axis=1)
        nearest = int(dists.argmin())
        if dists[nearest] < snap_radius:
            pv_new[vi] = cand_arr[nearest]

    return pv_new