"""Context-only correspondence proposals and fixed multi-view geometry certificates.""" from __future__ import annotations import copy import math from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Literal import torch from torch import Tensor, nn from mcss.types import Cameras @dataclass(frozen=True) class CertificateConfig: proposer: Literal["loftr_indoor_new"] = "loftr_indoor_new" confidence_threshold: float = 0.2 reciprocal_threshold_px: float = 1.5 cycle_threshold_px: float = 2.0 epipolar_threshold_px: float = 1.5 reprojection_threshold_px: float = 1.5 verification_threshold_px: float = 2.0 min_triangulation_angle_deg: float = 1.0 grid_stride: int = 8 device: str = "auto" model_cache_dir: str = "outputs/model_cache/torch" def __post_init__(self) -> None: if self.proposer != "loftr_indoor_new": raise ValueError("proposer must be loftr_indoor_new") if not 0.0 <= self.confidence_threshold <= 1.0: raise ValueError("confidence_threshold must be in [0, 1]") for name in ( "reciprocal_threshold_px", "cycle_threshold_px", "epipolar_threshold_px", "reprojection_threshold_px", "verification_threshold_px", "min_triangulation_angle_deg", ): if not math.isfinite(getattr(self, name)) or getattr(self, name) <= 0: raise ValueError(f"{name} must be positive and finite") if self.verification_threshold_px < self.reprojection_threshold_px: raise ValueError("verification threshold cannot be tighter than pair reprojection") if self.grid_stride < 1: raise ValueError("grid_stride must be positive") if not self.model_cache_dir: raise ValueError("model_cache_dir must be non-empty") @dataclass(frozen=True) class TriangulationResult: point: Tensor cheirality: bool reprojection_error_px: float triangulation_angle_deg: float @dataclass(frozen=True) class GeometryCertificate: source_view: int source_xy: tuple[float, float] matched_view: int matched_xy: tuple[float, float] verification_view: int verification_xy: tuple[float, float] point: tuple[float, float, float] descriptor_distance: float reciprocal_error_px: float cycle_error_px: float reprojection_error_px: float verification_error_px: float triangulation_angle_deg: float proposal_confidence: float = 1.0 epipolar_error_px: float = 0.0 @dataclass(frozen=True) class RawObservation: source_view: int source_xy: tuple[float, float] matched_view: int matched_xy: tuple[float, float] proposal_confidence: float point: tuple[float, float, float] | None cheirality: bool reprojection_error_px: float | None triangulation_angle_deg: float | None epipolar_error_px: float | None @dataclass(frozen=True) class CertificateBuildResult: image_size: tuple[int, int] context_views: int stage_counts: dict[str, int] certified_source_coverage: float source_grid_samples: int certificates: tuple[GeometryCertificate, ...] raw_observations: tuple[RawObservation, ...] pair_match_counts: dict[str, int] def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass(frozen=True) class _PairMatches: source: Tensor target: Tensor confidence: Tensor _MATCHER_CACHE: dict[tuple[str, str, str], nn.Module] = {} def fundamental_matrix(cameras: Cameras, source_view: int, target_view: int) -> Tensor: """Return the calibrated fundamental matrix mapping source points to target lines.""" _validate_unbatched_cameras(cameras) _validate_view_index(cameras, source_view) _validate_view_index(cameras, target_view) relative = torch.linalg.inv(cameras.c2w[target_view]) @ cameras.c2w[source_view] rotation = relative[:3, :3] translation = relative[:3, 3] tx = torch.stack( ( torch.stack((translation.new_zeros(()), -translation[2], translation[1])), torch.stack((translation[2], translation.new_zeros(()), -translation[0])), torch.stack((-translation[1], translation[0], translation.new_zeros(()))), ) ) essential = tx @ rotation matrix = ( torch.linalg.inv(cameras.intrinsics[target_view]).transpose(-1, -2) @ essential @ torch.linalg.inv(cameras.intrinsics[source_view]) ) return matrix / torch.linalg.vector_norm(matrix).clamp_min(torch.finfo(matrix.dtype).eps) def triangulate_dlt( source_xy: Tensor, target_xy: Tensor, cameras: Cameras, source_view: int, target_view: int, ) -> TriangulationResult: """Triangulate one calibrated correspondence with linear DLT.""" if source_xy.shape != (2,) or target_xy.shape != (2,): raise ValueError("source_xy and target_xy must have shape [2]") points, cheirality, reprojection, angle, finite = _triangulate_batch( source_xy.unsqueeze(0), target_xy.unsqueeze(0), cameras, source_view, target_view ) return TriangulationResult( point=points[0], cheirality=bool((cheirality & finite)[0].item()), reprojection_error_px=float(reprojection[0].item()), triangulation_angle_deg=float(angle[0].item()), ) def build_certificates( context_rgb: Tensor, context_cameras: Cameras, config: CertificateConfig ) -> CertificateBuildResult: """Build certificates without accepting any image target or geometry label.""" if context_rgb.ndim != 4 or context_rgb.shape[1] != 3: raise ValueError("context_rgb must have shape [V, 3, H, W]") views, _, height, width = context_rgb.shape if views < 3: raise ValueError("at least three context views are required for certification") if context_cameras.leading_shape != (views,): raise ValueError("context cameras must have one view dimension matching context_rgb") if context_cameras.image_size != (height, width): raise ValueError("camera image_size must match context_rgb") if not torch.isfinite(context_rgb).all(): raise ValueError("context_rgb must contain only finite values") device = _resolve_device(config.device) rgb = context_rgb.to(device=device, dtype=torch.float32) cameras = context_cameras.to(device=device, dtype=torch.float32) grayscale = ( rgb[:, 0:1] * 0.2989 + rgb[:, 1:2] * 0.5870 + rgb[:, 2:3] * 0.1140 ).clamp(0.0, 1.0) matcher = _get_matcher(config, device) pairs: dict[tuple[int, int], _PairMatches] = {} with torch.inference_mode(): for source_view in range(views): for target_view in range(views): if source_view == target_view: continue pairs[source_view, target_view] = _propose_pair( matcher, grayscale[source_view : source_view + 1], grayscale[target_view : target_view + 1], config.confidence_threshold, ) stage_counts = {name: 0 for name in ("raw", "reciprocal", "cycle", "geometric", "certified")} certificates: list[GeometryCertificate] = [] raw_observations: list[RawObservation] = [] pair_match_counts: dict[str, int] = {} certified_cells: set[tuple[int, int, int]] = set() for source_view in range(views): for matched_view in range(views): if source_view == matched_view: continue pair = pairs[source_view, matched_view] pair_match_counts[f"{source_view}->{matched_view}"] = len(pair.source) stage_counts["raw"] += len(pair.source) if len(pair.source) == 0: continue raw_points, raw_cheirality, raw_reprojection, raw_angle, raw_finite = ( _triangulate_batch( pair.source, pair.target, cameras, source_view, matched_view ) ) raw_epipolar = _epipolar_distance( pair.source, pair.target, fundamental_matrix(cameras, source_view, matched_view), ) for index in range(len(pair.source)): finite = bool(raw_finite[index].item()) point = _point_tuple(raw_points[index]) if finite else None raw_observations.append( RawObservation( source_view=source_view, source_xy=_xy_tuple(pair.source[index]), matched_view=matched_view, matched_xy=_xy_tuple(pair.target[index]), proposal_confidence=float(pair.confidence[index].item()), point=point, cheirality=bool((raw_cheirality & raw_finite)[index].item()), reprojection_error_px=( float(raw_reprojection[index].item()) if finite else None ), triangulation_angle_deg=float(raw_angle[index].item()) if finite else None, epipolar_error_px=float(raw_epipolar[index].item()) if finite else None, ) ) reverse = pairs[matched_view, source_view] reverse_index, reverse_source_error = _nearest(pair.target, reverse.source) reverse_target = _safe_gather(reverse.target, reverse_index, pair.source) reverse_target_error = torch.linalg.vector_norm(reverse_target - pair.source, dim=-1) reciprocal_error = torch.maximum(reverse_source_error, reverse_target_error) reciprocal = reciprocal_error <= config.reciprocal_threshold_px stage_counts["reciprocal"] += int(reciprocal.sum().item()) if not reciprocal.any(): continue original_indices = torch.nonzero(reciprocal, as_tuple=False)[:, 0] source_xy = pair.source[original_indices] matched_xy = pair.target[original_indices] confidence = pair.confidence[original_indices] reciprocal_error = reciprocal_error[original_indices] cycle_error, verification_view, verification_xy = _best_cycle( source_xy, matched_xy, pairs, source_view, matched_view, views, ) cycle = cycle_error <= config.cycle_threshold_px stage_counts["cycle"] += int(cycle.sum().item()) if not cycle.any(): continue source_xy = source_xy[cycle] matched_xy = matched_xy[cycle] confidence = confidence[cycle] reciprocal_error = reciprocal_error[cycle] cycle_error = cycle_error[cycle] verification_view = verification_view[cycle] verification_xy = verification_xy[cycle] points, cheirality, reprojection, angle, finite = _triangulate_batch( source_xy, matched_xy, cameras, source_view, matched_view ) epipolar = _epipolar_distance( source_xy, matched_xy, fundamental_matrix(cameras, source_view, matched_view), ) geometric = ( finite & cheirality & (epipolar <= config.epipolar_threshold_px) & (reprojection <= config.reprojection_threshold_px) & (angle >= config.min_triangulation_angle_deg) ) stage_counts["geometric"] += int(geometric.sum().item()) if not geometric.any(): continue source_xy = source_xy[geometric] matched_xy = matched_xy[geometric] confidence = confidence[geometric] reciprocal_error = reciprocal_error[geometric] cycle_error = cycle_error[geometric] verification_view = verification_view[geometric] verification_xy = verification_xy[geometric] points = points[geometric] reprojection = reprojection[geometric] angle = angle[geometric] epipolar = epipolar[geometric] verification_projection, verification_valid = _project_by_view( points, cameras, verification_view ) verification_error = torch.linalg.vector_norm( verification_projection - verification_xy, dim=-1 ) certified = verification_valid & ( verification_error <= config.verification_threshold_px ) stage_counts["certified"] += int(certified.sum().item()) for index in torch.nonzero(certified, as_tuple=False)[:, 0].tolist(): source_cell_x = min(width - 1, max(0, int(source_xy[index, 0].item()))) source_cell_y = min(height - 1, max(0, int(source_xy[index, 1].item()))) certified_cells.add( ( source_view, source_cell_y // config.grid_stride, source_cell_x // config.grid_stride, ) ) certificates.append( GeometryCertificate( source_view=source_view, source_xy=_xy_tuple(source_xy[index]), matched_view=matched_view, matched_xy=_xy_tuple(matched_xy[index]), verification_view=int(verification_view[index].item()), verification_xy=_xy_tuple(verification_xy[index]), point=_point_tuple(points[index]), descriptor_distance=float(1.0 - confidence[index].item()), reciprocal_error_px=float(reciprocal_error[index].item()), cycle_error_px=float(cycle_error[index].item()), reprojection_error_px=float(reprojection[index].item()), verification_error_px=float(verification_error[index].item()), triangulation_angle_deg=float(angle[index].item()), proposal_confidence=float(confidence[index].item()), epipolar_error_px=float(epipolar[index].item()), ) ) grid_height = math.ceil(height / config.grid_stride) grid_width = math.ceil(width / config.grid_stride) source_grid_samples = views * grid_height * grid_width coverage = len(certified_cells) / source_grid_samples certificates.sort( key=lambda item: ( item.source_view, item.source_xy[1], item.source_xy[0], item.matched_view, ) ) raw_observations.sort( key=lambda item: ( item.source_view, item.source_xy[1], item.source_xy[0], item.matched_view, ) ) return CertificateBuildResult( image_size=(height, width), context_views=views, stage_counts=stage_counts, certified_source_coverage=coverage, source_grid_samples=source_grid_samples, certificates=tuple(certificates), raw_observations=tuple(raw_observations), pair_match_counts=pair_match_counts, ) def _get_matcher(config: CertificateConfig, device: torch.device) -> nn.Module: cache_path = str(Path(config.model_cache_dir).resolve()) key = (config.proposer, str(device), cache_path) if key not in _MATCHER_CACHE: from kornia.feature import LoFTR from kornia.feature.loftr.loftr import default_cfg Path(cache_path).mkdir(parents=True, exist_ok=True) torch.hub.set_dir(cache_path) matcher = LoFTR(pretrained="indoor_new", config=copy.deepcopy(default_cfg)) matcher.requires_grad_(False).eval().to(device) _MATCHER_CACHE[key] = matcher return _MATCHER_CACHE[key] def _propose_pair( matcher: nn.Module, image_source: Tensor, image_target: Tensor, confidence_threshold: float ) -> _PairMatches: output = matcher({"image0": image_source, "image1": image_target}) source = output["keypoints0"].to(dtype=torch.float32) target = output["keypoints1"].to(dtype=torch.float32) confidence = output["confidence"].to(dtype=torch.float32) valid = ( torch.isfinite(source).all(dim=-1) & torch.isfinite(target).all(dim=-1) & torch.isfinite(confidence) & (confidence >= confidence_threshold) ) source = source[valid] target = target[valid] confidence = confidence[valid] if len(source): width = image_source.shape[-1] order = torch.argsort(source[:, 1] * (width + 1) + source[:, 0], stable=True) source = source[order] target = target[order] confidence = confidence[order] return _PairMatches(source, target, confidence) def _best_cycle( source_xy: Tensor, matched_xy: Tensor, pairs: dict[tuple[int, int], _PairMatches], source_view: int, matched_view: int, views: int, ) -> tuple[Tensor, Tensor, Tensor]: count = len(source_xy) best_error = torch.full((count,), torch.inf, device=source_xy.device) best_view = torch.full((count,), -1, dtype=torch.long, device=source_xy.device) best_xy = torch.zeros_like(source_xy) for third_view in range(views): if third_view in {source_view, matched_view}: continue forward = pairs[matched_view, third_view] backward = pairs[third_view, source_view] forward_index, forward_error = _nearest(matched_xy, forward.source) third_xy = _safe_gather(forward.target, forward_index, matched_xy) backward_index, backward_error = _nearest(third_xy, backward.source) returned_xy = _safe_gather(backward.target, backward_index, source_xy) closure_error = torch.linalg.vector_norm(returned_xy - source_xy, dim=-1) cycle_error = torch.maximum( torch.maximum(forward_error, backward_error), closure_error ) better = cycle_error < best_error best_error = torch.where(better, cycle_error, best_error) best_view = torch.where( better, torch.full_like(best_view, third_view), best_view ) best_xy = torch.where(better.unsqueeze(-1), third_xy, best_xy) return best_error, best_view, best_xy def _triangulate_batch( source_xy: Tensor, target_xy: Tensor, cameras: Cameras, source_view: int, target_view: int, ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: _validate_unbatched_cameras(cameras) _validate_view_index(cameras, source_view) _validate_view_index(cameras, target_view) if source_xy.ndim != 2 or source_xy.shape[-1] != 2 or target_xy.shape != source_xy.shape: raise ValueError("source and target pixels must share shape [N, 2]") projection_source = ( cameras.intrinsics[source_view] @ torch.linalg.inv(cameras.c2w[source_view])[:3] ) projection_target = ( cameras.intrinsics[target_view] @ torch.linalg.inv(cameras.c2w[target_view])[:3] ) rows = torch.stack( ( source_xy[:, 0, None] * projection_source[2] - projection_source[0], source_xy[:, 1, None] * projection_source[2] - projection_source[1], target_xy[:, 0, None] * projection_target[2] - projection_target[0], target_xy[:, 1, None] * projection_target[2] - projection_target[1], ), dim=1, ) _, _, vh = torch.linalg.svd(rows) homogeneous = vh[:, -1] epsilon = torch.finfo(homogeneous.dtype).eps finite = torch.isfinite(homogeneous).all(dim=-1) & (homogeneous[:, 3].abs() > epsilon) safe_scale = torch.where( homogeneous[:, 3].abs() > epsilon, homogeneous[:, 3], torch.ones_like(homogeneous[:, 3]) ) points = homogeneous[:, :3] / safe_scale[:, None] finite = finite & torch.isfinite(points).all(dim=-1) projected_source, depth_source = _project_points(points, cameras, source_view) projected_target, depth_target = _project_points(points, cameras, target_view) reprojection = torch.maximum( torch.linalg.vector_norm(projected_source - source_xy, dim=-1), torch.linalg.vector_norm(projected_target - target_xy, dim=-1), ) cheirality = (depth_source > epsilon) & (depth_target > epsilon) center_source = cameras.c2w[source_view, :3, 3] center_target = cameras.c2w[target_view, :3, 3] ray_source = points - center_source ray_target = points - center_target ray_source = ray_source / torch.linalg.vector_norm(ray_source, dim=-1, keepdim=True).clamp_min( epsilon ) ray_target = ray_target / torch.linalg.vector_norm(ray_target, dim=-1, keepdim=True).clamp_min( epsilon ) cosine = (ray_source * ray_target).sum(dim=-1).clamp(-1.0, 1.0) angle = torch.rad2deg(torch.acos(cosine)) return points, cheirality, reprojection, angle, finite def _project_points(points: Tensor, cameras: Cameras, view: int) -> tuple[Tensor, Tensor]: homogeneous = torch.cat((points, torch.ones_like(points[:, :1])), dim=-1) camera_points = homogeneous @ torch.linalg.inv(cameras.c2w[view])[:3].transpose(0, 1) projected = camera_points @ cameras.intrinsics[view].transpose(0, 1) epsilon = torch.finfo(points.dtype).eps safe_depth = torch.where( projected[:, 2].abs() > epsilon, projected[:, 2], torch.ones_like(projected[:, 2]) ) return projected[:, :2] / safe_depth[:, None], camera_points[:, 2] def _project_by_view(points: Tensor, cameras: Cameras, views: Tensor) -> tuple[Tensor, Tensor]: pixels = torch.zeros_like(points[:, :2]) valid = torch.zeros((len(points),), dtype=torch.bool, device=points.device) height, width = cameras.image_size for view in torch.unique(views).tolist(): if view < 0: continue mask = views == view projected, depth = _project_points(points[mask], cameras, int(view)) pixels[mask] = projected valid[mask] = ( torch.isfinite(projected).all(dim=-1) & (depth > torch.finfo(depth.dtype).eps) & (projected[:, 0] >= 0) & (projected[:, 0] <= width - 1) & (projected[:, 1] >= 0) & (projected[:, 1] <= height - 1) ) return pixels, valid def _epipolar_distance(source_xy: Tensor, target_xy: Tensor, matrix: Tensor) -> Tensor: homogeneous_source = torch.cat((source_xy, torch.ones_like(source_xy[:, :1])), dim=-1) homogeneous_target = torch.cat((target_xy, torch.ones_like(target_xy[:, :1])), dim=-1) target_lines = homogeneous_source @ matrix.transpose(0, 1) numerator = (homogeneous_target * target_lines).sum(dim=-1).abs() denominator = torch.linalg.vector_norm(target_lines[:, :2], dim=-1).clamp_min( torch.finfo(source_xy.dtype).eps ) return numerator / denominator def _nearest(query: Tensor, candidates: Tensor) -> tuple[Tensor, Tensor]: if len(candidates) == 0: return ( torch.zeros((len(query),), dtype=torch.long, device=query.device), torch.full((len(query),), torch.inf, device=query.device), ) indices = [] distances = [] for start in range(0, len(query), 1024): matrix = torch.cdist(query[start : start + 1024], candidates) distance, index = matrix.min(dim=1) indices.append(index) distances.append(distance) return torch.cat(indices), torch.cat(distances) def _safe_gather(values: Tensor, indices: Tensor, fallback: Tensor) -> Tensor: if len(values) == 0: return torch.full_like(fallback, torch.nan) return values[indices] def _resolve_device(requested: str) -> torch.device: if requested == "auto": return torch.device("cuda" if torch.cuda.is_available() else "cpu") return torch.device(requested) def _validate_unbatched_cameras(cameras: Cameras) -> None: if len(cameras.leading_shape) != 1: raise ValueError("certificate geometry requires unbatched cameras with shape [V, ...]") def _validate_view_index(cameras: Cameras, view: int) -> None: if not 0 <= view < cameras.leading_shape[0]: raise IndexError(f"camera view {view} is out of range") def _xy_tuple(value: Tensor) -> tuple[float, float]: return float(value[0].item()), float(value[1].item()) def _point_tuple(value: Tensor) -> tuple[float, float, float]: return float(value[0].item()), float(value[1].item()), float(value[2].item())