src/learned_baseline/postprocess_v2.py

"""Post-processing functions for segment predictions."""
import numpy as np


# ── Helpers ─────────────────────────────────────────────────────────────────

def _compact(pv, pe):
    """Remove unused vertices, remap edge indices."""
    if len(pe) == 0:
        return pv[:0], pe
    used, inv = np.unique(pe.reshape(-1), return_inverse=True)
    return pv[used], inv.reshape(-1, 2).astype(pe.dtype)


# ── New postprocessing (adapted from LundUni approach) ───────────────────────

def filter_short_edges(pv, pe, min_len=0.10):
    """Remove edges shorter than min_len metres (likely noise / merged artefacts)."""
    if len(pe) == 0:
        return pv, pe
    lens = np.linalg.norm(pv[pe[:, 1]] - pv[pe[:, 0]], axis=1)
    pe = pe[lens >= min_len]
    return _compact(pv, pe)


def remove_solitary_edges(pv, pe):
    """Remove edges where BOTH endpoints have degree == 1 (isolated noise edges)."""
    if len(pe) == 0:
        return pv, pe
    deg = np.bincount(pe.reshape(-1), minlength=len(pv))
    # keep edge if at least one endpoint is shared with another edge
    keep = np.any(deg[pe] > 1, axis=1)
    pe = pe[keep]
    return _compact(pv, pe)


def remove_near_duplicate_edges(pv, pe, dist_thresh=0.30):
    """Remove edges whose midpoint + direction are very close to a longer edge.

    Processes edges longest-first (like LundUni); removes shorter duplicates.
    """
    if len(pe) < 2:
        return pv, pe

    lengths = np.linalg.norm(pv[pe[:, 1]] - pv[pe[:, 0]], axis=1)
    order = np.argsort(-lengths)           # longest first
    pe_s = pe[order]
    lens_s = lengths[order]

    mids = 0.5 * (pv[pe_s[:, 0]] + pv[pe_s[:, 1]])
    dirs = pv[pe_s[:, 1]] - pv[pe_s[:, 0]]
    dirs = dirs / (lens_s[:, None] + 1e-8)

    keep = np.ones(len(pe_s), dtype=bool)
    for i in range(len(pe_s)):
        if not keep[i]:
            continue
        # Distance from later edges' midpoints to this edge's infinite line
        dp = mids[i + 1:] - mids[i]
        proj = (dp * dirs[i]).sum(axis=1, keepdims=True) * dirs[i]
        perp = np.linalg.norm(dp - proj, axis=1)
        # Also check directional similarity
        cos_sim = np.abs((dirs[i + 1:] * dirs[i]).sum(axis=1))
        suppress = (perp < dist_thresh) & (cos_sim > 0.90)
        keep[i + 1:][suppress] = False

    pe_out = pe_s[keep]
    # restore original index order (not strictly needed but cleaner)
    return _compact(pv, pe_out)


def reposition_to_line_intersections(pv, pe, max_move=0.50):
    """Move each vertex to the least-squares intersection of its incident lines.

    Replaces centroid-averaging with a proper line-intersection solve.
    Adapted from LundUni wireframe_postprocess.py concept.
    Only moves vertex if solution is within max_move metres of original position.
    """
    if len(pe) == 0:
        return pv
    new_pv = pv.copy()
    deg = np.bincount(pe.reshape(-1), minlength=len(pv))

    for vi in range(len(pv)):
        if deg[vi] < 2:
            continue
        # Collect incident edge directions
        mask = (pe[:, 0] == vi) | (pe[:, 1] == vi)
        inc = pe[mask]
        lines = []
        for a, b in inc:
            other = b if a == vi else a
            d = pv[other] - pv[vi]
            n = float(np.linalg.norm(d))
            if n < 1e-6:
                continue
            lines.append((pv[vi].copy(), d / n))

        if len(lines) < 2:
            continue

        # Build least-squares system: minimise sum ||(I - d d^T)(x - p)||^2
        # Normal equations: (sum (I - d d^T)) x = sum (I - d d^T) p
        A = np.zeros((3, 3), dtype=np.float64)
        b = np.zeros(3, dtype=np.float64)
        for p, d in lines:
            P = np.eye(3) - np.outer(d, d)
            A += P
            b += P @ p

        try:
            x, _, _, _ = np.linalg.lstsq(A, b, rcond=None)
            if np.linalg.norm(x - pv[vi]) < max_move:
                new_pv[vi] = x.astype(np.float32)
        except Exception:
            pass

    return new_pv


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