rpx_benchmark.adapters.batched_depth

  1"""True-batched depth model wrapper.
  2
  3The stock :class:`rpx_benchmark.adapters.base.BenchmarkableModel` iterates
  4samples one at a time inside ``predict``: even with ``--batch-size 8``,
  5the underlying adapter still sees one image per call, so frameworks that
  6batch internally (HuggingFace pipelines, DataLoader-driven PyTorch
  7inference) get no benefit. ``BatchedDepthBenchmarkModel`` fixes that for
  8monocular metric depth: it hands the whole list of RGBs to the adapter
  9in one call, lets the adapter do a single batched forward, then unwraps
 10the depth predictions one-per-sample.
 11
 12The model itself is untouched — only the dispatch shape changes — so
 13metric values are bit-identical to the per-sample path. Only throughput
 14improves (≈1.3× at batch 8 for ZoeDepth on an RTX 5070 Laptop;
 15typically larger for smaller models).
 16
 17Adapter contract
 18----------------
 19The wrapped adapter must accept either:
 20
 21* a single ``np.ndarray`` of shape ``(H, W, 3)`` uint8 → return depth ``(H', W')`` float
 22* a ``list[np.ndarray]`` → return ``list[np.ndarray]`` (one depth array per RGB,
 23  same order)
 24
 25If the adapter returns a depth shape different from its input RGB, we
 26bilinearly resize back to the RGB's HW so the depth aligns with the GT
 27mask (PIL-only, no OpenCV dep).
 28
 29Save layout
 30-----------
 31When ``save_dir`` is provided, predictions are written as
 32``<save_dir>/<scene>/<phase>/<frame>.npz`` (key ``"depth"``, compressed
 33float32 in the model's native output space) — the same scene/phase shape as
 34the on-disk dataset, so downstream analytics and per-scene plotting share one
 35navigation pattern. Writes are atomic and compression is outside inference
 36timing.
 37
 38Example
 39-------
 40::
 41
 42    from rpx_benchmark.adapters import BatchedDepthBenchmarkModel
 43    from depth_models.zoedepth import ZoeDepth
 44
 45    adapter = ZoeDepth(device="cuda", batch_size=8)
 46    model = BatchedDepthBenchmarkModel(
 47        adapter, name="ZoeDepth_NK",
 48        save_dir="rpx_results/ZoeDepth_NK/easy/predictions",
 49    )
 50    runner = BenchmarkRunner(model=model, dataset=dataset, ...)
 51"""
 52
 53from __future__ import annotations
 54
 55import os
 56import tempfile
 57import zipfile
 58from concurrent.futures import ThreadPoolExecutor
 59from pathlib import Path
 60from typing import Any, List, Optional, Sequence
 61
 62import numpy as np
 63from PIL import Image
 64
 65from ..api import DepthPrediction, Sample, TaskType
 66
 67__all__ = ["BatchedDepthBenchmarkModel"]
 68
 69
 70def _resize_bilinear_2d(src: np.ndarray, target_hw: tuple[int, int]) -> np.ndarray:
 71    """Resize a 2D float array to ``(H, W)``. PIL-only — no OpenCV dependency."""
 72    img = Image.fromarray(src.astype(np.float32), mode="F")
 73    img = img.resize((target_hw[1], target_hw[0]), Image.BILINEAR)
 74    return np.asarray(img, dtype=np.float32)
 75
 76
 77class BatchedDepthBenchmarkModel:
 78    """Depth model wrapper that batches predictions across a sample list.
 79
 80    Parameters
 81    ----------
 82    adapter
 83        Callable: ``adapter(rgb_list_or_single) -> depth_list_or_single``.
 84        Single-image: ``np.ndarray (H, W, 3) uint8`` → ``np.ndarray (H', W') float``.
 85        Batched: ``list[np.ndarray]`` → ``list[np.ndarray]``.
 86    name
 87        Display name surfaced in result.json and Box paths.
 88    save_dir
 89        If provided, writes per-frame ``.npz`` predictions under
 90        ``<save_dir>/<scene>/<phase>/<frame>.npz``. Created lazily on the
 91        first save (no side effect at construction time).
 92    native_alignment
 93        Declared by the adapter (``"none"`` for metric models like ZoeDepth /
 94        UniDepth V2 / Depth Pro; ``"ls_affine"`` for up-to-scale models
 95        like Marigold / Lotus-2). The runner reads this to apply the
 96        right alignment by default. Falls back to ``"none"`` if the
 97        adapter doesn't declare it. The wrapper returns raw predictions;
 98        the runner applies this mode once per complete scene/phase cell.
 99    native_precision
100        Declared by the adapter (``"fp32"``, ``"fp16"``, ``"bf16"``).
101        Used by the runner to record the OperatingPoint precision
102        tag. Falls back to ``"fp32"`` if not declared.
103    allow_nonpositive_predictions
104        Permit finite, non-degenerate raw depth maps containing zero or
105        negative values. This is opt-in for protocols that preserve raw model
106        output and apply a documented evaluation-domain transform later.
107    """
108
109    task = TaskType.MONOCULAR_DEPTH
110
111    def __init__(
112        self,
113        adapter: Any,
114        *,
115        name: str,
116        save_dir: Optional[str | Path] = None,
117        native_alignment: Optional[str] = None,
118        native_precision: Optional[str] = None,
119        resume_predictions: bool = False,
120        allow_nonpositive_predictions: bool = False,
121    ) -> None:
122        self._adapter = adapter
123        self.name = name
124        self._save_dir = Path(save_dir) if save_dir else None
125        if resume_predictions and self._save_dir is None:
126            from ..exceptions import ConfigError
127
128            raise ConfigError("resume_predictions requires save_dir")
129        self.resume_predictions = bool(resume_predictions)
130        self.allow_nonpositive_predictions = bool(allow_nonpositive_predictions)
131        self.resume_stats = {
132            "cache_hits": 0,
133            "inferred_new": 0,
134            "invalid_recomputed": 0,
135        }
136        self.last_cache_hits: list[bool] = []
137        self._last_save_required: list[bool] = []
138        # Read alignment off the adapter if it declared one, else fall
139        # back to the constructor arg, else "none" (metric assumption).
140        self.native_alignment: str = (
141            native_alignment or getattr(adapter, "native_alignment", None) or "none"
142        )
143        self.native_precision: str = (
144            native_precision or getattr(adapter, "native_precision", None) or "fp32"
145        )
146        self.depth_output_kind: str = "relative" if self.native_alignment != "none" else "metric"
147        # Profiler walker reaches the underlying nn.Module via this attr.
148        self.model = adapter
149
150    # No-op: the adapter loads its own weights at construction.
151    def setup(self) -> None:
152        return None
153
154    def predict(self, batch: Sequence[Sample]) -> List[DepthPrediction]:
155        cached: dict[int, np.ndarray] = {}
156        missing_indices: list[int] = []
157        invalid_existing: set[int] = set()
158        for index, sample in enumerate(batch):
159            path = self._prediction_path(sample)
160            if self.resume_predictions and path is not None and path.exists():
161                depth = self._load_valid_prediction(
162                    path,
163                    np.asarray(sample.rgb).shape[:2],
164                    require_positive=not self.allow_nonpositive_predictions,
165                )
166                if depth is not None:
167                    cached[index] = depth
168                    self.resume_stats["cache_hits"] += 1
169                    continue
170                invalid_existing.add(index)
171            missing_indices.append(index)
172
173        inferred: list[np.ndarray] = []
174        if missing_indices:
175            rgbs = [np.asarray(batch[i].rgb, dtype=np.uint8) for i in missing_indices]
176            raw_depths = self._adapter(rgbs)
177            if not isinstance(raw_depths, (list, tuple)):
178                raw_depths = [raw_depths]
179            inferred = [np.asarray(depth, dtype=np.float32) for depth in raw_depths]
180        if len(inferred) != len(missing_indices):
181            from ..exceptions import AdapterError
182
183            raise AdapterError(
184                f"adapter returned {len(inferred)} depths for {len(missing_indices)} inputs",
185                hint="Check the adapter's batched contract: it must return "
186                "one depth map per input RGB.",
187            )
188
189        inferred_by_index = dict(zip(missing_indices, inferred, strict=True))
190        preds: List[DepthPrediction] = []
191        self.last_cache_hits = []
192        self._last_save_required = []
193        for index, sample in enumerate(batch):
194            is_cached = index in cached
195            d_raw = cached[index] if is_cached else inferred_by_index[index]
196            target_hw = np.asarray(sample.rgb).shape[:2]
197            if d_raw.shape != target_hw:
198                d_raw = _resize_bilinear_2d(d_raw, target_hw)
199            if not is_cached:
200                if index in invalid_existing:
201                    self.resume_stats["invalid_recomputed"] += 1
202                else:
203                    self.resume_stats["inferred_new"] += 1
204
205            # Return the raw prediction. Relative outputs are deliberately
206            # aligned later by BenchmarkRunner, after every frame in the
207            # same (scene, phase) cell is available for one pooled solve.
208            preds.append(DepthPrediction(depth_map=d_raw))
209            self.last_cache_hits.append(is_cached)
210            self._last_save_required.append(not is_cached)
211        return preds
212
213    def persist_predictions(
214        self,
215        batch: Sequence[Sample],
216        predictions: Sequence[DepthPrediction],
217    ) -> None:
218        """Persist newly inferred predictions after the runner stops timing."""
219        pending = [
220            (sample, prediction.depth_map)
221            for sample, prediction, required in zip(
222                batch,
223                predictions,
224                self._last_save_required,
225                strict=True,
226            )
227            if required
228        ]
229        if len(pending) <= 1:
230            for sample, depth in pending:
231                self._maybe_save(sample, depth)
232            return
233
234        # Compression is CPU-bound and each frame has a distinct atomic target.
235        # Parallelize it within the batch so a fast GPU is not followed by a
236        # serial chain of np.savez_compressed calls.
237        workers = min(len(pending), 8)
238        with ThreadPoolExecutor(max_workers=workers) as executor:
239            futures = [
240                executor.submit(self._maybe_save, sample, depth)
241                for sample, depth in pending
242            ]
243            for future in futures:
244                future.result()
245
246    def _prediction_path(self, sample: Sample) -> Path | None:
247        if self._save_dir is None:
248            return None
249        # Prefer the manifest's metadata block (set by local_manifest.py
250        # and dataset_hub.split_manifests so the loader rides it through
251        # to Sample.metadata). Fall back to id-parsing.
252        meta = getattr(sample, "metadata", None) or {}
253        scene = meta.get("scene_id")
254        phase = meta.get("phase_idx")
255        frame = meta.get("frame")
256        if not (scene and frame is not None and phase is not None):
257            try:
258                scene, phase, frame = str(sample.id).split("__", 2)
259            except ValueError:
260                scene, phase, frame = "unknown", "0", str(sample.id)
261        return self._save_dir / str(scene) / str(phase) / f"{frame}.npz"
262
263    @staticmethod
264    def _load_valid_prediction(
265        path: Path,
266        target_hw: tuple[int, int],
267        *,
268        require_positive: bool = True,
269    ) -> np.ndarray | None:
270        try:
271            with np.load(path, allow_pickle=False) as payload:
272                if payload.files != ["depth"]:
273                    return None
274                depth = np.asarray(payload["depth"])
275            return BatchedDepthBenchmarkModel._validated_depth(
276                depth,
277                target_hw,
278                require_positive=require_positive,
279            )
280        except (OSError, ValueError, KeyError, EOFError, zipfile.BadZipFile):
281            return None
282
283    @staticmethod
284    def _validated_depth(
285        depth: np.ndarray,
286        target_hw: tuple[int, int],
287        *,
288        require_positive: bool = True,
289    ) -> np.ndarray | None:
290        depth = np.asarray(depth)
291        if depth.ndim != 2 or depth.shape != target_hw:
292            return None
293        if depth.dtype != np.dtype(np.float32):
294            return None
295        if not np.isfinite(depth).all():
296            return None
297        if require_positive and np.any(depth <= 0):
298            return None
299        if float(np.ptp(depth)) <= 1e-6:
300            return None
301        return depth
302
303    def _maybe_save(self, sample: Sample, depth: np.ndarray) -> None:
304        out = self._prediction_path(sample)
305        if out is None:
306            return
307        canonical = np.asarray(depth, dtype=np.float32)
308        target_hw = np.asarray(sample.rgb).shape[:2]
309        if (
310            self._validated_depth(
311                canonical,
312                target_hw,
313                require_positive=not self.allow_nonpositive_predictions,
314            )
315            is None
316        ):
317            from ..exceptions import AdapterError
318
319            raise AdapterError(
320                f"Refusing to save an invalid depth prediction for sample {sample.id!r}: "
321                f"expected a finite, non-degenerate float32 map with shape {target_hw}"
322                + (
323                    " and strictly positive values."
324                    if not self.allow_nonpositive_predictions
325                    else "."
326                )
327            )
328        out.parent.mkdir(parents=True, exist_ok=True)
329        temp_path: Path | None = None
330        try:
331            with tempfile.NamedTemporaryFile(
332                mode="w+b",
333                suffix=".npz.part",
334                dir=out.parent,
335                delete=False,
336            ) as handle:
337                temp_path = Path(handle.name)
338                np.savez_compressed(handle, depth=canonical)
339                handle.flush()
340                os.fsync(handle.fileno())
341            os.replace(temp_path, out)
342        finally:
343            if temp_path is not None and temp_path.exists():
344                temp_path.unlink()