"""Known-camera consistency certificates for dense supplier depth.""" from __future__ import annotations import hashlib import json import os import shutil import tempfile from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Self import numpy as np from mcss.geometry_supplier import GeometrySupplierResult DENSE_CERTIFICATE_CACHE_SCHEMA_VERSION = 1 @dataclass(frozen=True) class DenseCertificateConfig: grid_stride: int = 4 min_depth_m: float = 0.10 max_depth_m: float = 20.0 depth_absolute_tolerance_m: float = 0.10 depth_relative_tolerance: float = 0.05 roundtrip_threshold_px: float = 2.0 surface_half_width_m: float = 0.05 free_space_margin_m: float = 0.05 minimum_supporting_views: int = 1 maximum_conflicting_views: int = 0 def __post_init__(self) -> None: if self.grid_stride < 1: raise ValueError("grid_stride must be positive") if not 0 < self.min_depth_m < self.max_depth_m: raise ValueError("depth range must be positive and increasing") values = ( self.depth_absolute_tolerance_m, self.depth_relative_tolerance, self.roundtrip_threshold_px, self.surface_half_width_m, self.free_space_margin_m, ) if min(values) < 0: raise ValueError("certificate tolerances must be nonnegative") if self.minimum_supporting_views < 1 or self.maximum_conflicting_views < 0: raise ValueError("support/conflict counts are invalid") @dataclass class DenseGeometryCertificates: frame_ids: tuple[int, ...] pixel_xy: np.ndarray depth_along_ray: np.ndarray world_points: np.ndarray raw_confidence: np.ndarray raw_valid: np.ndarray certified: np.ndarray support_count: np.ndarray conflict_count: np.ndarray occlusion_count: np.ndarray provenance: np.ndarray free_end_m: np.ndarray surface_near_m: np.ndarray surface_far_m: np.ndarray risk: np.ndarray def __post_init__(self) -> None: scalar_shape = self.depth_along_ray.shape if len(scalar_shape) != 3 or scalar_shape[0] != len(self.frame_ids): raise ValueError("dense certificate scalar fields require shape [V, Gh, Gw]") scalar_fields = ( self.raw_confidence, self.raw_valid, self.certified, self.support_count, self.conflict_count, self.occlusion_count, self.free_end_m, self.surface_near_m, self.surface_far_m, self.risk, ) if any(value.shape != scalar_shape for value in scalar_fields): raise ValueError("dense certificate scalar field shapes must match") if self.pixel_xy.shape != (*scalar_shape, 2): raise ValueError("pixel_xy must have shape [V, Gh, Gw, 2]") if self.world_points.shape != (*scalar_shape, 3): raise ValueError("world_points must have shape [V, Gh, Gw, 3]") if self.provenance.shape != (*scalar_shape, len(self.frame_ids)): raise ValueError("provenance must have one trailing source-view dimension") for value in ( self.pixel_xy, self.depth_along_ray, self.world_points, self.raw_confidence, self.free_end_m, self.surface_near_m, self.surface_far_m, self.risk, ): if not np.isfinite(value).all(): raise ValueError("dense certificate numeric fields must be finite") if np.any((self.risk < 0) | (self.risk > 1)): raise ValueError("risk must be in [0, 1]") def canonical(self) -> Self: order = np.argsort(np.asarray(self.frame_ids), kind="stable") if np.array_equal(order, np.arange(len(order))): return self provenance = self.provenance[order][..., order] values = { name: getattr(self, name)[order] for name in ( "pixel_xy", "depth_along_ray", "world_points", "raw_confidence", "raw_valid", "certified", "support_count", "conflict_count", "occlusion_count", "free_end_m", "surface_near_m", "surface_far_m", "risk", ) } return type(self)( tuple(self.frame_ids[int(index)] for index in order), provenance=provenance, **values, ) def build_dense_certificates( supplier: GeometrySupplierResult, config: DenseCertificateConfig ) -> DenseGeometryCertificates: supplier = supplier.canonical() view_count, height, width = supplier.depth_along_ray.shape y_values = np.arange(0, height, config.grid_stride, dtype=np.int64) x_values = np.arange(0, width, config.grid_stride, dtype=np.int64) grid_y, grid_x = np.meshgrid(y_values, x_values, indexing="ij") grid_shape = grid_y.shape pixel_grid = np.stack((grid_x, grid_y), axis=-1).astype(np.float32) pixel_xy = np.broadcast_to(pixel_grid, (view_count, *grid_shape, 2)).copy() depth = supplier.depth_along_ray[:, grid_y, grid_x] confidence = supplier.confidence[:, grid_y, grid_x] raw_mask = supplier.valid_mask[:, grid_y, grid_x] raw_valid = ( raw_mask & np.isfinite(depth) & (depth >= config.min_depth_m) & (depth <= config.max_depth_m) ) world_points = np.zeros((view_count, *grid_shape, 3), dtype=np.float32) directions = [] for view in range(view_count): camera_directions = _camera_directions(pixel_grid, supplier.intrinsics[view]) directions.append(camera_directions) world_directions = camera_directions @ supplier.c2w[view, :3, :3].T world_points[view] = ( supplier.c2w[view, :3, 3] + world_directions * depth[view][..., None] ) support = np.zeros((view_count, *grid_shape), dtype=np.int16) conflict = np.zeros_like(support) occlusion = np.zeros_like(support) provenance = np.zeros((view_count, *grid_shape, view_count), dtype=np.bool_) residual_sum = np.zeros((view_count, *grid_shape), dtype=np.float32) residual_count = np.zeros_like(support) for source in range(view_count): provenance[source, ..., source] = raw_valid[source] for target in range(view_count): if target == source: continue projected_xy, projected_distance, in_front = _project_world( world_points[source], supplier.intrinsics[target], supplier.c2w[target] ) inside = ( in_front & (projected_xy[..., 0] >= 0) & (projected_xy[..., 0] <= width - 1) & (projected_xy[..., 1] >= 0) & (projected_xy[..., 1] <= height - 1) ) target_depth, target_valid = _sample_nearest( supplier.depth_along_ray[target], supplier.valid_mask[target], projected_xy, ) comparable = ( raw_valid[source] & inside & target_valid & np.isfinite(target_depth) & (target_depth >= config.min_depth_m) & (target_depth <= config.max_depth_m) ) tolerance = ( config.depth_absolute_tolerance_m + config.depth_relative_tolerance * projected_distance ) difference = target_depth - projected_distance target_world = _world_from_pixels( projected_xy, target_depth, supplier.intrinsics[target], supplier.c2w[target] ) roundtrip_xy, _, roundtrip_front = _project_world( target_world, supplier.intrinsics[source], supplier.c2w[source] ) roundtrip_error = np.linalg.norm(roundtrip_xy - pixel_grid, axis=-1) agrees = ( comparable & (np.abs(difference) <= tolerance) & roundtrip_front & (roundtrip_error <= config.roundtrip_threshold_px) ) is_occlusion = comparable & (difference < -tolerance) is_conflict = comparable & ~agrees & ~is_occlusion support[source] += agrees.astype(np.int16) occlusion[source] += is_occlusion.astype(np.int16) conflict[source] += is_conflict.astype(np.int16) provenance[source, ..., target] = agrees normalized_residual = np.minimum( 1.0, np.abs(difference) / np.maximum(tolerance, 1e-8), ) residual_sum[source] += np.where(comparable, normalized_residual, 0.0) residual_count[source] += comparable.astype(np.int16) certified = ( raw_valid & (support >= config.minimum_supporting_views) & (conflict <= config.maximum_conflicting_views) ) free_end = np.where(certified, np.maximum(0.0, depth - config.free_space_margin_m), 0.0) surface_near = np.where(certified, np.maximum(0.0, depth - config.surface_half_width_m), 0.0) surface_far = np.where(certified, depth + config.surface_half_width_m, 0.0) possible_support = max(1, view_count - 1) support_deficit = 1.0 - np.minimum(support, possible_support) / possible_support conflict_rate = conflict / possible_support residual = residual_sum / np.maximum(residual_count, 1) confidence_uncertainty = 1.0 / (1.0 + np.maximum(confidence, 0.0)) risk = np.clip( 0.35 * support_deficit + 0.30 * residual + 0.20 * conflict_rate + 0.15 * confidence_uncertainty, 0.0, 1.0, ).astype(np.float32) return DenseGeometryCertificates( frame_ids=supplier.frame_ids, pixel_xy=pixel_xy, depth_along_ray=depth.astype(np.float32), world_points=world_points, raw_confidence=confidence.astype(np.float32), raw_valid=raw_valid, certified=certified, support_count=support, conflict_count=conflict, occlusion_count=occlusion, provenance=provenance, free_end_m=free_end.astype(np.float32), surface_near_m=surface_near.astype(np.float32), surface_far_m=surface_far.astype(np.float32), risk=risk, ) def write_dense_certificate_cache( destination: str | Path, certificates: DenseGeometryCertificates, config: DenseCertificateConfig, *, supplier_arrays_sha256: str, ) -> Path: """Seal dense certificates in an overwrite-refusing JSON+NPZ directory.""" target = Path(destination).resolve() if target.exists(): raise FileExistsError(f"refusing to overwrite dense certificate cache: {target}") if not _is_sha256(supplier_arrays_sha256): raise ValueError("supplier_arrays_sha256 must be a lowercase SHA-256 digest") target.parent.mkdir(parents=True, exist_ok=True) temporary = Path( tempfile.mkdtemp(prefix=f".{target.name}.tmp-", dir=str(target.parent)) ) canonical = certificates.canonical() try: arrays_path = temporary / "arrays.npz" arrays = _certificate_arrays(canonical) np.savez_compressed(arrays_path, **arrays) metadata = { "schema_version": DENSE_CERTIFICATE_CACHE_SCHEMA_VERSION, "frame_ids": list(canonical.frame_ids), "config": asdict(config), "supplier_arrays_sha256": supplier_arrays_sha256, "arrays_file": arrays_path.name, "arrays_sha256": _sha256_file(arrays_path), "arrays": { name: {"shape": list(value.shape), "dtype": str(value.dtype)} for name, value in arrays.items() }, } (temporary / "metadata.json").write_text( json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8" ) os.replace(temporary, target) except BaseException: shutil.rmtree(temporary, ignore_errors=True) raise return target def read_dense_certificate_cache( source: str | Path, ) -> tuple[DenseGeometryCertificates, dict[str, Any]]: root = Path(source).resolve() metadata_path = root / "metadata.json" if not metadata_path.is_file(): raise FileNotFoundError(metadata_path) metadata = json.loads(metadata_path.read_text(encoding="utf-8")) if ( not isinstance(metadata, dict) or metadata.get("schema_version") != DENSE_CERTIFICATE_CACHE_SCHEMA_VERSION ): raise ValueError("dense certificate cache metadata has an unsupported schema") arrays_path = root / str(metadata.get("arrays_file", "")) if not arrays_path.is_file(): raise FileNotFoundError(arrays_path) if _sha256_file(arrays_path) != metadata.get("arrays_sha256"): raise ValueError("dense certificate cache arrays hash does not match metadata") with np.load(arrays_path, allow_pickle=False) as values: required = set(_certificate_array_names()) | {"frame_ids"} if set(values.files) != required: raise ValueError("dense certificate cache array keys do not match schema") frame_ids = tuple(int(value) for value in values["frame_ids"].tolist()) certificates = DenseGeometryCertificates( frame_ids=frame_ids, **{name: values[name].copy() for name in _certificate_array_names()}, ) if list(frame_ids) != metadata.get("frame_ids"): raise ValueError("dense certificate metadata frame IDs do not match arrays") actual_schema = { name: {"shape": list(value.shape), "dtype": str(value.dtype)} for name, value in _certificate_arrays(certificates).items() } if actual_schema != metadata.get("arrays"): raise ValueError("dense certificate metadata array schema does not match arrays") return certificates, metadata def _camera_directions(pixel_xy: np.ndarray, intrinsics: np.ndarray) -> np.ndarray: homogeneous = np.concatenate( (pixel_xy, np.ones((*pixel_xy.shape[:-1], 1), dtype=np.float32)), axis=-1 ) directions = homogeneous @ np.linalg.inv(intrinsics).T return directions / np.maximum(np.linalg.norm(directions, axis=-1, keepdims=True), 1e-8) def _world_from_pixels( pixel_xy: np.ndarray, depth: np.ndarray, intrinsics: np.ndarray, c2w: np.ndarray ) -> np.ndarray: camera_direction = _camera_directions(pixel_xy, intrinsics) world_direction = camera_direction @ c2w[:3, :3].T return c2w[:3, 3] + world_direction * depth[..., None] def _project_world( world: np.ndarray, intrinsics: np.ndarray, c2w: np.ndarray ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: rotation = c2w[:3, :3] camera = (world - c2w[:3, 3]) @ rotation projected = camera @ intrinsics.T z = projected[..., 2] xy = projected[..., :2] / np.maximum(z[..., None], 1e-8) distance = np.linalg.norm(camera, axis=-1) return xy.astype(np.float32), distance.astype(np.float32), camera[..., 2] > 1e-8 def _sample_nearest( depth: np.ndarray, mask: np.ndarray, pixel_xy: np.ndarray ) -> tuple[np.ndarray, np.ndarray]: height, width = depth.shape x = np.clip(np.rint(pixel_xy[..., 0]).astype(np.int64), 0, width - 1) y = np.clip(np.rint(pixel_xy[..., 1]).astype(np.int64), 0, height - 1) return depth[y, x], mask[y, x] def _certificate_array_names() -> tuple[str, ...]: return ( "pixel_xy", "depth_along_ray", "world_points", "raw_confidence", "raw_valid", "certified", "support_count", "conflict_count", "occlusion_count", "provenance", "free_end_m", "surface_near_m", "surface_far_m", "risk", ) def _certificate_arrays(certificates: DenseGeometryCertificates) -> dict[str, np.ndarray]: result = { name: np.asarray(getattr(certificates, name)) for name in _certificate_array_names() } result["frame_ids"] = np.asarray(certificates.frame_ids, dtype=np.int64) return result def _sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for block in iter(lambda: handle.read(1024 * 1024), b""): digest.update(block) return digest.hexdigest() def _is_sha256(value: str) -> bool: return len(value) == 64 and all(character in "0123456789abcdef" for character in value)