colmap_refine.py

"""COLMAP-based vertex position refinement.

Two complementary refinement strategies that use the COLMAP sparse point
cloud as a high-precision 3D landmark source:

1. ``refine_vertices_3d_plane`` — Variant (a+c).
   For each merged_v vertex, find its K nearest COLMAP points in 3D,
   fit a local plane, and project the vertex onto that plane. Cancels
   depth-noise residuals after the initial unprojection.

2. ``refine_vertices_multiview_plane`` — Variant (b).
   For each merged_v vertex, project it into every view, find the K
   nearest COLMAP points in 2D within each view's image, fit a local
   plane in 3D from those points, project the vertex onto the plane,
   and average the resulting 3D positions across views weighted by the
   plane fit quality.

Both methods only use ``pycolmap`` + ``numpy`` + ``scipy``. Purely
geometric — no thresholds tuned on local validation.
"""

from __future__ import annotations

import numpy as np
from scipy.spatial import cKDTree

from hoho2025.example_solutions import convert_entry_to_human_readable

try:
    from mvs_utils import collect_views, project_world_to_image
except ImportError:
    from submission.mvs_utils import collect_views, project_world_to_image


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _fit_plane_pca(points: np.ndarray) -> tuple[np.ndarray, np.ndarray, float]:
    """PCA plane fit. Returns (centroid, unit_normal, fit_quality).

    fit_quality = 1 - (smallest_eigval / largest_eigval). 1.0 = perfectly
    planar, 0.0 = sphere. Used as a weight when combining multi-view
    refinements.
    """
    centroid = points.mean(axis=0)
    centred = points - centroid
    # SVD instead of eig to be numerically stable on small N
    _, s, Vt = np.linalg.svd(centred, full_matrices=False)
    if len(s) < 3:
        return centroid, np.array([0.0, 1.0, 0.0]), 0.0
    normal = Vt[2]  # smallest variance direction
    # quality: ratio of last to first singular value, inverted
    if s[0] < 1e-9:
        return centroid, normal, 0.0
    quality = 1.0 - float(s[2] / s[0])
    return centroid, normal, max(0.0, min(1.0, quality))


def _project_point_to_plane(
    point: np.ndarray, plane_centroid: np.ndarray, plane_normal: np.ndarray,
) -> np.ndarray:
    """Orthogonal projection of ``point`` onto a plane defined by
    ``(centroid, unit normal)``.
    """
    rel = point - plane_centroid
    d = float(np.dot(rel, plane_normal))
    return point - d * plane_normal


# ---------------------------------------------------------------------------
# Variant (a+c): 3D KD-tree neighbours → local plane → snap
# ---------------------------------------------------------------------------

def refine_vertices_3d_plane(
    vertices: np.ndarray,
    colmap_xyz: np.ndarray,
    knn_radius: float = 0.5,
    knn_k: int = 12,
    min_neighbours: int = 6,
    max_displacement: float = 0.5,
    min_quality: float = 0.6,
) -> tuple[np.ndarray, np.ndarray]:
    """Refine each vertex by snapping to a local plane fit through its
    nearest COLMAP neighbours in 3D.

    Parameters
    ----------
    vertices : (N, 3) array of merged 3D vertex positions.
    colmap_xyz : (M, 3) all COLMAP points3D world coordinates.
    knn_radius : maximum distance for a neighbour to count.
    knn_k : maximum number of neighbours to use (for speed).
    min_neighbours : refuse to refine when fewer neighbours found.
    max_displacement : reject the snap if it moves the vertex by more
        than this many metres (likely a wall plane, not the roof).
    min_quality : reject when the local plane fit is not flat enough
        (PCA quality below this).

    Returns
    -------
    refined : (N, 3) refined vertex positions.
    snapped : (N,)  bool — which vertices were moved.
    """
    verts = np.asarray(vertices, dtype=np.float64)
    refined = verts.copy()
    snapped = np.zeros(len(verts), dtype=bool)

    if len(verts) == 0 or len(colmap_xyz) < min_neighbours:
        return refined, snapped

    tree = cKDTree(colmap_xyz)
    for i, v in enumerate(verts):
        idx = tree.query_ball_point(v, knn_radius)
        if len(idx) < min_neighbours:
            continue
        if len(idx) > knn_k:
            # Pick the closest knn_k of the candidates
            d = np.linalg.norm(colmap_xyz[idx] - v, axis=1)
            order = np.argsort(d)[:knn_k]
            idx = [idx[j] for j in order]

        nbrs = colmap_xyz[idx]
        centroid, normal, quality = _fit_plane_pca(nbrs)
        if quality < min_quality:
            continue

        projected = _project_point_to_plane(v, centroid, normal)
        if float(np.linalg.norm(projected - v)) > max_displacement:
            continue
        refined[i] = projected
        snapped[i] = True

    return refined, snapped


# ---------------------------------------------------------------------------
# Variant (b): multi-view consensus plane refinement
# ---------------------------------------------------------------------------

def refine_vertices_multiview_plane(
    vertices: np.ndarray,
    entry,
    knn_2d_px: float = 30.0,
    knn_k: int = 12,
    min_neighbours: int = 6,
    max_displacement: float = 0.5,
    min_quality: float = 0.5,
    min_views: int = 2,
) -> tuple[np.ndarray, np.ndarray]:
    """Multi-view consensus refinement.

    For each vertex:
      1. Project it into every available view.
      2. In each view, find COLMAP points whose own 2D projection is
         within ``knn_2d_px`` of the vertex projection.
      3. Take the corresponding 3D points and fit a local plane.
      4. Project the vertex onto that plane → one candidate 3D position
         per view, weighted by the plane's PCA quality.
      5. Combine the per-view candidates as a quality-weighted mean.

    Crucially, the 2D pixel neighbourhood ensures the COLMAP points used
    for the plane fit are the **ones the camera sees near this vertex** —
    not just close in 3D — so it does not blend roof + wall + ground
    points like a 3D KNN would.

    Returns ``(refined, snapped)`` arrays in the same shape as the input.
    """
    verts = np.asarray(vertices, dtype=np.float64)
    refined = verts.copy()
    snapped = np.zeros(len(verts), dtype=bool)

    if len(verts) == 0:
        return refined, snapped

    good = convert_entry_to_human_readable(entry)
    colmap_rec = good.get('colmap') or good.get('colmap_binary')
    if colmap_rec is None:
        return refined, snapped

    views = collect_views(colmap_rec, good['image_ids'])
    if len(views) < 1:
        return refined, snapped

    colmap_xyz = np.array(
        [p.xyz for p in colmap_rec.points3D.values()], dtype=np.float64
    )
    if len(colmap_xyz) < min_neighbours:
        return refined, snapped

    # Pre-project all COLMAP points into each view once
    per_view_proj: dict[str, tuple[np.ndarray, np.ndarray]] = {}
    for vid, info in views.items():
        uv, z = project_world_to_image(info['P'], colmap_xyz)
        in_front = z > 0
        per_view_proj[vid] = (uv[in_front], np.where(in_front)[0])

    for i, v in enumerate(verts):
        candidates: list[tuple[np.ndarray, float]] = []
        for vid, info in views.items():
            uv_v, z_v = project_world_to_image(info['P'], v.reshape(1, 3))
            if z_v[0] <= 0:
                continue
            target_uv = uv_v[0]
            H, W = info['height'], info['width']
            if not (0 <= target_uv[0] < W and 0 <= target_uv[1] < H):
                continue
            view_uv, view_idx = per_view_proj[vid]
            if len(view_uv) == 0:
                continue
            d = np.linalg.norm(view_uv - target_uv, axis=1)
            mask = d <= knn_2d_px
            if mask.sum() < min_neighbours:
                continue
            cand_idx = view_idx[mask]
            d_in = d[mask]
            if len(cand_idx) > knn_k:
                order = np.argsort(d_in)[:knn_k]
                cand_idx = cand_idx[order]
            nbrs = colmap_xyz[cand_idx]
            centroid, normal, quality = _fit_plane_pca(nbrs)
            if quality < min_quality:
                continue
            projected = _project_point_to_plane(v, centroid, normal)
            if float(np.linalg.norm(projected - v)) > max_displacement:
                continue
            candidates.append((projected, quality))

        if len(candidates) < min_views:
            continue

        # Quality-weighted mean
        weights = np.array([c[1] for c in candidates], dtype=np.float64)
        positions = np.array([c[0] for c in candidates], dtype=np.float64)
        if weights.sum() < 1e-6:
            continue
        new_pos = (positions * weights[:, None]).sum(axis=0) / weights.sum()
        refined[i] = new_pos
        snapped[i] = True

    return refined, snapped