rpx_benchmark.metrics.depth_alignment

  1"""Per-frame alignment of a predicted depth map to GT scale.
  2
  3Used by both:
  4
  5* the runner's primary-metric path (``BatchedDepthBenchmarkModel.predict``
  6  applies the adapter's ``native_alignment`` before the runner sees the
  7  prediction, so result.json's ``aggregated.absrel`` is meaningful for
  8  relative-depth models like Marigold / MiDaS); and
  9* the post-processor's metric basket (``comprehensive_depth_metrics``)
 10  which iterates alignment modes per-frame.
 11
 12Single source of truth for the alignment math so a fix in one path
 13applies everywhere. Modes match the standard mono-depth-evaluation
 14literature.
 15"""
 16
 17from __future__ import annotations
 18
 19from typing import Optional
 20
 21import numpy as np
 22
 23from ..exceptions import ConfigError
 24
 25__all__ = [
 26    "DEPTH_MIN_M",
 27    "DEPTH_MAX_M",
 28    "ALIGNMENT_MODES",
 29    "align_pred_to_gt",
 30    "align_pred_to_gt_pooled",
 31    "default_valid_mask",
 32]
 33
 34
 35#: D435 sensor working range. Pixels outside this range are excluded
 36#: from the alignment fit (and from every metric — see
 37#: ``comprehensive_depth_metrics._valid``). Kept here so any future
 38#: change is one-place.
 39DEPTH_MIN_M: float = 0.3
 40DEPTH_MAX_M: float = 5.0
 41
 42#: Supported alignment mode names.
 43ALIGNMENT_MODES = frozenset({"none", "median", "ls_affine", "ls_disparity", "ls_log"})
 44
 45
 46def default_valid_mask(pred: np.ndarray, gt_m: np.ndarray) -> np.ndarray:
 47    """Standard validity gate used by every alignment + metric step.
 48
 49    A pixel enters the alignment fit (and the downstream metric) iff:
 50
 51    * GT is finite (D435 holes encoded as NaN or 0 mm are excluded);
 52    * GT is in the sensor working range (``DEPTH_MIN_M`` … ``DEPTH_MAX_M``);
 53    * Prediction is finite;
 54    * Prediction is positive.
 55    """
 56    return (
 57        np.isfinite(gt_m)
 58        & (gt_m > DEPTH_MIN_M)
 59        & (gt_m < DEPTH_MAX_M)
 60        & np.isfinite(pred)
 61        & (pred > 0)
 62    )
 63
 64
 65def align_pred_to_gt(
 66    pred: np.ndarray,
 67    gt_m: np.ndarray,
 68    mode: str,
 69    valid: Optional[np.ndarray] = None,
 70) -> np.ndarray:
 71    """Fit a per-frame transform from raw prediction to GT scale.
 72
 73    Parameters
 74    ----------
 75    pred
 76        Predicted depth (any 2-D float array; semantics depend on the
 77        model — for relative models it's up-to-scale, for metric models
 78        it's already in metres).
 79    gt_m
 80        Ground-truth depth in metres.
 81    mode
 82        One of:
 83
 84        * ``"none"``         — identity. Use for native-metric models.
 85        * ``"median"``       — ``pred *= median(gt) / median(pred)``;
 86          cheap rescale, single dof.
 87        * ``"ls_affine"``    — least-squares fit ``a·pred + b`` to GT
 88          in depth space. Two dof; standard alignment for relative
 89          and up-to-scale models (DA-V2 relative, MiDaS, Marigold,
 90          Lotus).
 91        * ``"ls_disparity"`` — least-squares fit of a model's native
 92          inverse-depth/disparity output to GT disparity, followed by
 93          conversion to metric depth.
 94        * ``"ls_log"``       — least-squares fit ``a*pred + b`` to
 95          ``log(GT)``, then exponentiate. This is FE2E's published
 96          ``norm_type=ln`` output convention.
 97    valid
 98        Optional pre-computed validity mask. If not provided, uses
 99        :func:`default_valid_mask`.
100
101    Returns
102    -------
103    np.ndarray
104        Aligned prediction with the same shape and dtype as ``pred``.
105
106    Raises
107    ------
108    ConfigError
109        If ``mode`` isn't in :data:`ALIGNMENT_MODES`.
110    """
111    if mode == "none" or pred is None or gt_m is None:
112        return pred
113    if mode not in ALIGNMENT_MODES:
114        raise ConfigError(
115            f"unknown alignment mode {mode!r}",
116            hint=f"Expected one of: {sorted(ALIGNMENT_MODES)}.",
117        )
118    if valid is None:
119        valid = default_valid_mask(pred, gt_m)
120    if not valid.any():
121        return pred
122
123    p = pred[valid].astype(np.float64)
124    g = gt_m[valid].astype(np.float64)
125
126    if mode == "median":
127        scale = float(np.median(g) / max(np.median(p), 1e-9))
128        return (pred * scale).astype(np.float32)
129
130    if mode == "ls_affine":
131        A = np.column_stack([p, np.ones_like(p)])
132        coef, *_ = np.linalg.lstsq(A, g, rcond=None)
133        a, b = float(coef[0]), float(coef[1])
134        return (a * pred + b).astype(np.float32)
135
136    if mode == "ls_disparity":
137        eps = 1e-3
138        inv_g = 1.0 / np.maximum(g, eps)
139        A = np.column_stack([p, np.ones_like(p)])
140        coef, *_ = np.linalg.lstsq(A, inv_g, rcond=None)
141        a, b = float(coef[0]), float(coef[1])
142        aligned_disparity = a * pred.astype(np.float64) + b
143        return (1.0 / np.maximum(aligned_disparity, eps)).astype(np.float32)
144
145    if mode == "ls_log":
146        A = np.column_stack([p, np.ones_like(p)])
147        coef, *_ = np.linalg.lstsq(A, np.log(g), rcond=None)
148        a, b = float(coef[0]), float(coef[1])
149        log_depth = np.minimum(a * pred.astype(np.float64) + b, 5.0)
150        return np.exp(log_depth).astype(np.float32)
151
152    # Unreachable thanks to the ALIGNMENT_MODES gate above; raises
153    # defensively in case the gate is loosened later.
154    raise ConfigError(
155        f"unknown alignment mode {mode!r}",
156        hint="Add the mode to ALIGNMENT_MODES and implement its branch.",
157    )
158
159
160def align_pred_to_gt_pooled(
161    pred_seq: np.ndarray,
162    gt_seq: np.ndarray,
163    mode: str,
164    valid_seq: Optional[np.ndarray] = None,
165) -> np.ndarray:
166    """Pooled scale+shift fit over an entire ``(scene, phase)`` cell.
167
168    Paper §3.3 (line 62) prescribes **per-scene** scale-and-shift
169    alignment for affine-invariant depth models (Lotus-2, FE2E, etc.).
170    "Per-scene" here means one ``(s, t)`` per ``(scene, phase)`` cell,
171    pooling every valid pixel across all frames in that cell into a
172    single least-squares solve. The same scope applies to:
173
174    * Video Depth (paper task ``video_depth``) — pool over ``(T, H, W)``.
175    * Image Depth (paper task ``monocular_depth``) — pool over the phase's
176      ~250 frames, with each ``(H, W)`` slice contributing its valid
177      pixels.
178
179    Per-frame alignment would let each frame find its own ``(s, t)``,
180    artificially smoothing over any temporal scale drift the model has
181    within a phase — hiding a real failure mode. Per-scene preserves
182    that drift and reports it through the per-frame metrics (which are
183    then averaged into the cell).
184
185    Parameters
186    ----------
187    pred_seq
188        Predicted depth for the cell. Either ``(T, H, W)`` (a phase
189        clip) or a list of per-frame arrays stacked along axis 0. Any
190        2-D-shape-compatible array works; only the first axis is
191        treated as the frame index for masking purposes.
192    gt_seq
193        Ground-truth depth in metres, same shape as ``pred_seq``.
194    mode
195        One of :data:`ALIGNMENT_MODES`. ``"none"`` returns ``pred_seq``
196        unchanged. ``"median"`` / ``"ls_affine"`` / ``"ls_disparity"``
197        all pool valid pixels across frames before solving.
198    valid_seq
199        Optional pre-computed mask, same shape as ``pred_seq``. If
200        absent, computed per-frame with :func:`default_valid_mask` and
201        OR'd to form the pooled mask.
202
203    Returns
204    -------
205    np.ndarray
206        Aligned prediction with the same shape and dtype as
207        ``pred_seq``. When the pooled valid set is empty, returns
208        ``pred_seq`` unchanged.
209
210    Notes
211    -----
212    A native-metric model gets ``mode="none"`` (or any mode where the
213    solver finds ``s ≈ 1, t ≈ 0`` because the prediction already
214    matches the GT distribution). An affine-invariant model gets
215    ``mode="ls_affine"``. The runner is expected to dispatch on the
216    adapter's ``depth_output_kind`` attribute.
217    """
218    if mode == "none":
219        return pred_seq
220    if mode not in ALIGNMENT_MODES:
221        raise ConfigError(
222            f"unknown alignment mode {mode!r}",
223            hint=f"Expected one of: {sorted(ALIGNMENT_MODES)}.",
224        )
225    if pred_seq.shape != gt_seq.shape:
226        raise ConfigError(
227            f"align_pred_to_gt_pooled: pred_seq.shape={pred_seq.shape} "
228            f"does not match gt_seq.shape={gt_seq.shape}",
229            hint="Both inputs must be the same shape (typically (T, H, W)).",
230        )
231    if valid_seq is None:
232        # Compute per-frame masks and stack — equivalent to computing
233        # one mask over the whole pooled array, but reuses the existing
234        # default_valid_mask policy without making it sequence-aware.
235        valid_seq = default_valid_mask(pred_seq, gt_seq)
236    if valid_seq.shape != pred_seq.shape:
237        raise ConfigError(
238            f"align_pred_to_gt_pooled: valid_seq.shape={valid_seq.shape} "
239            f"does not match pred_seq.shape={pred_seq.shape}",
240            hint="The validity mask must be the same shape as the depth tensors.",
241        )
242    if not valid_seq.any():
243        return pred_seq
244
245    # Flatten across (T, H, W) and let the existing per-frame solver
246    # do the math. We tag the prediction's dtype so the caller's
247    # contract (float32 / float64) is preserved.
248    p = pred_seq[valid_seq].astype(np.float64)
249    g = gt_seq[valid_seq].astype(np.float64)
250    if p.size == 0:
251        return pred_seq
252
253    if mode == "median":
254        scale = float(np.median(g) / max(np.median(p), 1e-9))
255        return (pred_seq * scale).astype(pred_seq.dtype)
256
257    if mode == "ls_affine":
258        A = np.column_stack([p, np.ones_like(p)])
259        coef, *_ = np.linalg.lstsq(A, g, rcond=None)
260        a, b = float(coef[0]), float(coef[1])
261        return (a * pred_seq + b).astype(pred_seq.dtype)
262
263    if mode == "ls_disparity":
264        eps = 1e-3
265        inv_g = 1.0 / np.maximum(g, eps)
266        A = np.column_stack([p, np.ones_like(p)])
267        coef, *_ = np.linalg.lstsq(A, inv_g, rcond=None)
268        a, b = float(coef[0]), float(coef[1])
269        aligned_disparity = a * pred_seq.astype(np.float64) + b
270        return (1.0 / np.maximum(aligned_disparity, eps)).astype(pred_seq.dtype)
271
272    if mode == "ls_log":
273        A = np.column_stack([p, np.ones_like(p)])
274        coef, *_ = np.linalg.lstsq(A, np.log(g), rcond=None)
275        a, b = float(coef[0]), float(coef[1])
276        log_depth = np.minimum(a * pred_seq.astype(np.float64) + b, 5.0)
277        return np.exp(log_depth).astype(pred_seq.dtype)
278
279    raise ConfigError(  # pragma: no cover — guarded by ALIGNMENT_MODES gate
280        f"unknown alignment mode {mode!r}",
281        hint="Add the mode to ALIGNMENT_MODES and implement its branch.",
282    )