src/mcss/v8_suppliers.py
"""Adapters for the frozen V8 geometry-supplier qualification branch."""
from __future__ import annotations
import contextlib
import importlib
import sys
from pathlib import Path
from typing import Any
import numpy as np
import torch
from mcss.geometry_supplier import GeometrySupplierRequest
from mcss.types import Cameras
from mcss.v8_geometry import (
V8Proposal,
depth_z_to_points_cam,
emvsnet_mm_to_m,
)
EMVSNET_DEPTH_NEAR_M = 0.10
EMVSNET_DEPTH_FAR_M = 20.0
EMVSNET_METRES_TO_CHECKPOINT_UNITS = 1000.0
def center_emvsnet_images(images: torch.Tensor, *, epsilon: float = 1e-8) -> torch.Tensor:
"""Apply the released EMVSNet per-image, per-channel center normalization."""
if images.ndim != 4 or images.shape[1] != 3:
raise ValueError("EMVSNet images must have shape [V, 3, H, W]")
if epsilon <= 0:
raise ValueError("epsilon must be positive")
values = images.to(dtype=torch.float32)
mean = values.mean(dim=(-2, -1), keepdim=True)
variance = values.var(dim=(-2, -1), correction=0, keepdim=True)
return (values - mean) / (torch.sqrt(variance) + epsilon)
def emvsnet_projection_coordinates(
cameras: Cameras, *, translation_scale: float
) -> torch.Tensor:
"""Build EMVSNet P=K[R|t] matrices in an explicitly scaled metric space."""
if not np.isfinite(translation_scale) or translation_scale <= 0:
raise ValueError("translation_scale must be positive and finite")
c2w = cameras.c2w.clone()
c2w[..., :3, 3] *= float(translation_scale)
w2c = torch.linalg.inv(c2w)
projection = torch.zeros(
(*cameras.leading_shape, 4, 4),
dtype=cameras.dtype,
device=cameras.device,
)
projection[..., 3, 3] = 1.0
projection[..., :3, :4] = torch.matmul(
cameras.intrinsics, w2c[..., :3, :4]
)
return projection
def emvsnet_epistemic_risk(
evidential: torch.Tensor, *, method: str
) -> torch.Tensor:
"""Compute the released EMVSNet epistemic uncertainty as an uncalibrated rank score."""
if evidential.ndim != 4 or evidential.shape[1] != 4:
raise ValueError("EMVSNet evidential output must have shape [B, 4, H, W]")
_gamma, nu, alpha, beta = torch.unbind(evidential, dim=1)
epsilon = torch.finfo(evidential.dtype).eps
if method == "der":
aleatoric = beta / (alpha - 1.0 + epsilon)
return aleatoric / (nu + epsilon)
if method == "sder":
return torch.rsqrt(nu + epsilon)
raise ValueError("method must be der or sder")
class UniDepthV2SmallSupplier:
"""Run UniDepthV2-Small independently for every context view."""
def __init__(
self,
*,
code_root: str | Path,
model_id: str = "lpiccinelli/unidepth-v2-vits14",
revision: str | None = None,
local_model_dir: str | Path | None = None,
model_cache_dir: str | Path | None = None,
device: str = "cuda",
) -> None:
self.code_root = Path(code_root).resolve()
self.model_id = model_id
self.revision = revision
self.local_model_dir = None if local_model_dir is None else Path(local_model_dir).resolve()
self.model_cache_dir = None if model_cache_dir is None else Path(model_cache_dir).resolve()
self.device = torch.device(device)
self._model: Any | None = None
def _load(self) -> Any:
if self._model is not None:
return self._model
if not self.code_root.is_dir():
raise FileNotFoundError(self.code_root)
with _prepend_path(self.code_root):
from unidepth.models import UniDepthV2
kwargs: dict[str, Any] = {}
model_source = self.model_id
if self.local_model_dir is not None:
if not (self.local_model_dir / "config.json").is_file() or not (
self.local_model_dir / "model.safetensors"
).is_file():
raise FileNotFoundError(
"local UniDepth bundle must contain config.json and model.safetensors"
)
model_source = str(self.local_model_dir)
kwargs["local_files_only"] = True
elif self.revision is not None:
kwargs["revision"] = self.revision
if self.model_cache_dir is not None:
# Hugging Face reads this environment variable during model resolution.
import os
os.environ.setdefault("HF_HOME", str(self.model_cache_dir))
self._model = UniDepthV2.from_pretrained(model_source, **kwargs)
self._model = self._model.to(self.device).eval()
for parameter in self._model.parameters():
parameter.requires_grad_(False)
return self._model
@torch.no_grad()
def infer(self, request: GeometrySupplierRequest) -> V8Proposal:
model = self._load()
canonical = request.canonical()
depths: list[np.ndarray] = []
points: list[np.ndarray] = []
confidences: list[np.ndarray] = []
valid: list[np.ndarray] = []
with torch.inference_mode():
for index in range(len(canonical.frame_ids)):
# UniDepth's public infer API expects uint8 RGB and performs its own normalization.
rgb = (canonical.rgb[index].clamp(0, 1) * 255.0).round().to(torch.uint8)
K = canonical.cameras.intrinsics[index].to(dtype=torch.float32)
output = model.infer(rgb, K)
depth = _squeeze_map(output, "depth")
pts = _squeeze_points(output, "points")
conf = _squeeze_map(output, "confidence")
if depth.shape != tuple(canonical.rgb.shape[-2:]):
raise ValueError("UniDepth output spatial size does not match request")
depths.append(depth)
points.append(pts)
confidences.append(conf)
valid.append(np.isfinite(depth) & (depth > 0))
return V8Proposal(
supplier="unidepth-v2-small",
frame_ids=canonical.frame_ids,
depth_z_m=np.stack(depths),
points_cam_m=np.stack(points),
native_confidence=np.stack(confidences),
native_risk=1.0 - np.clip(np.stack(confidences), 0.0, 1.0),
valid_mask=np.stack(valid),
intrinsics=canonical.cameras.intrinsics.detach().cpu().numpy(),
c2w=canonical.cameras.c2w.detach().cpu().numpy(),
uncertainty_provenance="unidepth_confidence_rank_only",
uncertainty_formula="model_output_confidence_not_calibrated",
)
class EMVSNetSupplier:
"""Run the pinned EMVSNet 64-bin model with an explicit mm->m boundary."""
def __init__(
self,
*,
code_root: str | Path,
checkpoint: str | Path,
device: str = "cuda",
image_scale: float = 0.25,
depth_bins: int = 64,
max_h: int = 512,
max_w: int = 640,
evidential_method: str = "der",
) -> None:
if depth_bins != 64:
raise ValueError("Q0 freezes the 64-bin EMVSNet checkpoint")
if evidential_method not in {"der", "sder"}:
raise ValueError("evidential_method must be der or sder")
self.code_root = Path(code_root).resolve()
self.checkpoint = Path(checkpoint).resolve()
self.device = torch.device(device)
self.image_scale = float(image_scale)
self.depth_bins = int(depth_bins)
self.max_h = int(max_h)
self.max_w = int(max_w)
self.evidential_method = evidential_method
self._model: Any | None = None
def depth_hypotheses_m(self, *, device: torch.device | None = None) -> torch.Tensor:
"""Return the pre-registered, label-independent Q0 metric depth range."""
return torch.linspace(
EMVSNET_DEPTH_NEAR_M,
EMVSNET_DEPTH_FAR_M,
self.depth_bins,
dtype=torch.float32,
device=device or self.device,
)
def _load(self) -> Any:
if self._model is not None:
return self._model
if not self.checkpoint.is_file():
raise FileNotFoundError(self.checkpoint)
with _prepend_path(self.code_root):
module = importlib.import_module("models.drmvsnet")
model = module.EMVSNet(
disparity_level=self.depth_bins,
image_scale=self.image_scale,
max_h=self.max_h,
max_w=self.max_w,
return_depth=False,
evidential_method=self.evidential_method,
)
state = torch.load(self.checkpoint, map_location="cpu", weights_only=False)
if not isinstance(state, dict) or "model" not in state:
raise ValueError("EMVSNet checkpoint must contain a model state_dict")
state_dict = {str(k).removeprefix("module."): v for k, v in state["model"].items()}
model.load_state_dict(state_dict, strict=True)
self._model = model.to(self.device).eval()
for parameter in self._model.parameters():
parameter.requires_grad_(False)
return self._model
@torch.no_grad()
def infer(self, request: GeometrySupplierRequest) -> V8Proposal:
if len(request.frame_ids) < 2:
raise ValueError("EMVSNet requires one reference and at least one source view")
model = self._load()
canonical = request.canonical()
height, width = canonical.cameras.image_size
expected_h = int(round(self.max_h * self.image_scale))
expected_w = int(round(self.max_w * self.image_scale))
if (height, width) != (expected_h, expected_w):
raise ValueError(
f"EMVSNet Q0 expects {(expected_h, expected_w)} input pixels, got {(height, width)}"
)
scaled_h, scaled_w = height, width
images = torch.nn.functional.interpolate(
canonical.rgb.to(self.device) * 255.0,
size=(scaled_h, scaled_w),
mode="bilinear",
align_corners=True,
)
images = center_emvsnet_images(images)
# EMVSNet's published DTU interface uses projection matrices P=K[R|t]. We construct
# those matrices from the OpenCV c2w cameras and feed a metric-equivalent mm scene.
cameras = canonical.cameras.to(self.device)
projection = emvsnet_projection_coordinates(
cameras, translation_scale=EMVSNET_METRES_TO_CHECKPOINT_UNITS
).unsqueeze(0)
depth_values_m = self.depth_hypotheses_m(device=self.device)
depth_values_mm = depth_values_m * EMVSNET_METRES_TO_CHECKPOINT_UNITS
with torch.inference_mode():
probability_volume, evidential, _branch_probabilities = model(
images.unsqueeze(0), projection, depth_values_mm.unsqueeze(0)
)
# EMVSNet's train/validation path regresses depth from the primary probability volume.
# The third output is the evidential branch ensemble and is not the reported depth head.
depth_mm = torch.sum(
probability_volume * depth_values_mm.view(1, -1, 1, 1), dim=1
)
depth_mm = torch.nn.functional.interpolate(
depth_mm[:, None], size=(height, width), mode="bilinear", align_corners=True
)[:, 0]
risk = torch.nn.functional.interpolate(
emvsnet_epistemic_risk(evidential, method=self.evidential_method)[:, None],
size=(height, width),
mode="bilinear",
align_corners=True,
)[:, 0]
depth_m = emvsnet_mm_to_m(depth_mm[0].detach().cpu().numpy())
# EMVSNet predicts one reference map. Replicate only as an explicit reference proposal;
# source views remain unknown rather than being fabricated.
depth_maps = np.zeros((len(canonical.frame_ids), height, width), dtype=np.float32)
points_maps = np.zeros((len(canonical.frame_ids), height, width, 3), dtype=np.float32)
conf_maps = np.zeros_like(depth_maps)
risk_maps = np.zeros_like(depth_maps)
valid_maps = np.zeros_like(depth_maps, dtype=np.bool_)
depth_maps[0] = depth_m
points_maps[0] = depth_z_to_points_cam(
depth_m, canonical.cameras.intrinsics[0].detach().cpu().numpy()
)
risk_maps[0] = risk[0].detach().cpu().numpy()
conf_maps[0] = 1.0 / (1.0 + risk_maps[0])
valid_maps[0] = np.isfinite(depth_m) & (depth_m > 0)
return V8Proposal(
supplier="emvsnet-64",
frame_ids=canonical.frame_ids,
depth_z_m=depth_maps,
points_cam_m=points_maps,
native_confidence=conf_maps,
native_risk=risk_maps,
valid_mask=valid_maps,
intrinsics=canonical.cameras.intrinsics.detach().cpu().numpy(),
c2w=canonical.cameras.c2w.detach().cpu().numpy(),
uncertainty_provenance="emvsnet_evidential_rank_only",
uncertainty_formula=f"{self.evidential_method}_publisher_release_inference_formula",
)
def _squeeze_map(output: dict[str, torch.Tensor], key: str) -> np.ndarray:
if key not in output:
raise ValueError(f"UniDepth output missing {key}")
value = output[key].detach().float().cpu()
while value.ndim > 2 and value.shape[0] == 1:
value = value[0]
if value.ndim == 3 and value.shape[0] == 1:
value = value[0]
if value.ndim != 2 or not torch.isfinite(value).all():
raise ValueError(f"UniDepth {key} has an invalid shape or nonfinite values")
return value.numpy()
def _squeeze_points(output: dict[str, torch.Tensor], key: str) -> np.ndarray:
if key not in output:
raise ValueError(f"UniDepth output missing {key}")
value = output[key].detach().float().cpu()
if value.ndim == 4 and value.shape[0] == 1:
value = value[0]
if value.ndim != 3 or value.shape[0] != 3:
raise ValueError("UniDepth points must have shape [1,3,H,W] or [3,H,W]")
if not torch.isfinite(value).all():
raise ValueError("UniDepth points must be finite")
return value.permute(1, 2, 0).numpy()
@contextlib.contextmanager
def _prepend_path(root: Path):
root_text = str(root)
old = list(sys.path)
if root_text not in sys.path:
sys.path.insert(0, root_text)
try:
yield
finally:
sys.path[:] = old