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