rpx_benchmark.adapters.batched_multimodal

  1"""Generic batched-dispatch BenchmarkModel base for multi-modal tasks.
  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 sample per call. ``BatchedDepthBenchmarkModel``
  6fixes that for monocular depth specifically.
  7
  8This module generalises the pattern so every other task (segmentation,
  9relative pose, sparse depth, keypoint matching, ...) can opt into
 10true-batched dispatch without each one having to re-implement the
 11batch-extraction loop, the per-sample save logic, and the prediction-
 12wrapping ceremony.
 13
 14Three task-specific hooks
 15-------------------------
 16Subclasses override:
 17
 18* :meth:`extract_inputs(sample) -> Any` — pull whatever the model needs
 19  off the sample (rgb, rgb_b, fisheye_left, intrinsics, …). The return
 20  value is opaque to the base class; the adapter receives a list of
 21  these.
 22* :meth:`call_adapter(inputs_list) -> outputs_list` — invoke the
 23  adapter on the full batch in one call. Default: ``self._adapter(inputs_list)``.
 24* :meth:`wrap_output(model_output, sample) -> Prediction` — wrap one
 25  model output into the task's Prediction dataclass (e.g.,
 26  ``DepthPrediction``, ``SegmentationPrediction``,
 27  ``RelativePosePrediction``).
 28
 29Optional:
 30
 31* :meth:`maybe_save(sample, model_output)` — persist per-frame outputs
 32  for offline analytics (the depth wrapper writes
 33  ``<save_dir>/<scene>/<phase>/<frame>.npz``; segmentation could write
 34  PNG masks; tracking writes JSON tracklets).
 35
 36Sample.metadata is the multi-modal side-channel
 37-----------------------------------------------
 38The loader's ``_load_sample`` puts non-RGB modalities into
 39``sample.metadata`` (``rgb_b``, ``fisheye_left``, ``fisheye_right``,
 40…). Subclasses pull from there:
 41
 42::
 43
 44    class BatchedRelativePoseBenchmarkModel(BatchedTaskBenchmarkModel):
 45        task = TaskType.RELATIVE_CAMERA_POSE
 46        def extract_inputs(self, sample):
 47            rgb_a = np.asarray(sample.rgb, dtype=np.uint8)
 48            rgb_b = np.asarray(sample.metadata["rgb_b"], dtype=np.uint8)
 49            return {"rgb_a": rgb_a, "rgb_b": rgb_b}
 50
 51The contract is intentionally loose: ``extract_inputs`` returns
 52whatever shape the adapter expects (numpy array, dict, tuple, ...),
 53and the adapter dispatches the batch as a list of those.
 54"""
 55
 56from __future__ import annotations
 57
 58from pathlib import Path
 59from typing import Any, List, Optional, Sequence
 60
 61import numpy as np
 62
 63from ..api import Sample, TaskType
 64
 65__all__ = [
 66    "BatchedTaskBenchmarkModel",
 67    "BatchedSegmentationBenchmarkModel",
 68    "BatchedRelativePoseBenchmarkModel",
 69]
 70
 71
 72class BatchedTaskBenchmarkModel:
 73    """Abstract base — true-batched-dispatch wrapper for any RPX task.
 74
 75    Subclasses set the class-level ``task`` (a :class:`TaskType`) and
 76    override ``extract_inputs`` + ``wrap_output``. The harness contract
 77    matches ``BenchmarkableModel`` (``setup`` + ``predict(batch)``).
 78
 79    Parameters
 80    ----------
 81    adapter
 82        The model callable. Receives a list of whatever ``extract_inputs``
 83        returns (one per sample). Returns a list of model outputs (one per
 84        sample, same order).
 85    name
 86        Display name surfaced in result.json + Box paths.
 87    save_dir
 88        If set, ``maybe_save`` writes per-frame artefacts. Default: no save.
 89    """
 90
 91    #: Task this wrapper produces predictions for. Subclasses set this.
 92    task: TaskType = TaskType.MONOCULAR_DEPTH  # placeholder — must be overridden
 93
 94    def __init__(
 95        self,
 96        adapter: Any,
 97        *,
 98        name: str,
 99        save_dir: Optional[str | Path] = None,
100    ) -> None:
101        self._adapter = adapter
102        self.name = name
103        self._save_dir = Path(save_dir) if save_dir else None
104        # Profiler walker reaches the underlying nn.Module via this attr.
105        self.model = adapter
106
107    def setup(self) -> None:
108        """No-op — the adapter loads its own weights at construction."""
109        return None
110
111    # ────────────────────────  hooks  ────────────────────────
112
113    def extract_inputs(self, sample: Sample) -> Any:
114        """Subclass: return whatever the adapter expects per sample.
115
116        Examples:
117        - Depth: ``np.asarray(sample.rgb, dtype=np.uint8)``
118        - Pose:  ``{"rgb_a": ..., "rgb_b": sample.metadata["rgb_b"]}``
119        - Stereo depth: ``{"left": metadata["fisheye_left"],
120                            "right": metadata["fisheye_right"]}``
121        """
122        raise NotImplementedError
123
124    def call_adapter(self, inputs: List[Any]) -> List[Any]:
125        """Subclass override point: how the adapter is invoked on a list.
126
127        Default: ``self._adapter(inputs)`` — works when the adapter
128        accepts a list directly (HF pipelines, our depth adapters).
129        Override for adapters with a different batched contract.
130        """
131        out = self._adapter(inputs)
132        if not isinstance(out, (list, tuple)):
133            out = [out]
134        return list(out)
135
136    def wrap_output(self, model_output: Any, sample: Sample) -> Any:
137        """Subclass: wrap one model output as the task's Prediction dataclass."""
138        raise NotImplementedError
139
140    def maybe_save(self, sample: Sample, model_output: Any) -> None:
141        """Optional override — persist per-frame artefacts for analytics."""
142        return None
143
144    # ────────────────────────  predict()  ────────────────────────
145
146    def predict(self, batch: Sequence[Sample]) -> List[Any]:
147        inputs = [self.extract_inputs(s) for s in batch]
148        outputs = self.call_adapter(inputs)
149        if len(outputs) != len(batch):
150            from ..exceptions import AdapterError
151
152            raise AdapterError(
153                f"adapter returned {len(outputs)} outputs for a batch of {len(batch)}",
154                hint="Check the adapter's batched contract: it must return "
155                "one output per input sample.",
156            )
157        preds: List[Any] = []
158        for sample, out in zip(batch, outputs, strict=False):
159            preds.append(self.wrap_output(out, sample))
160            self.maybe_save(sample, out)
161        return preds
162
163
164# ────────────────────────  Reference subclasses  ────────────────────────
165
166
167class BatchedSegmentationBenchmarkModel(BatchedTaskBenchmarkModel):
168    """Batched wrapper for object-segmentation models.
169
170    Adapter contract: ``adapter(list_of_rgb_uint8_HxWx3) -> list_of_int_mask_HxW``.
171    Output mask values are integer instance IDs (0 = background). If the
172    model returns a different shape than the input RGB, we nearest-neighbour
173    resize back to preserve integer IDs.
174    """
175
176    task = TaskType.OBJECT_SEGMENTATION
177
178    def extract_inputs(self, sample: Sample) -> np.ndarray:
179        return np.asarray(sample.rgb, dtype=np.uint8)
180
181    def wrap_output(self, model_output: Any, sample: Sample) -> Any:
182        from PIL import Image
183
184        from ..api import SegmentationPrediction
185
186        mask = np.asarray(model_output)
187        target_hw = np.asarray(sample.rgb).shape[:2]
188        if mask.shape != target_hw:
189            # Nearest-neighbour preserves integer IDs.
190            img = Image.fromarray(mask.astype(np.int32), mode="I")
191            img = img.resize((target_hw[1], target_hw[0]), Image.NEAREST)
192            mask = np.asarray(img, dtype=np.int32)
193        return SegmentationPrediction(mask=mask.astype(np.int32))
194
195    def maybe_save(self, sample: Sample, model_output: Any) -> None:
196        if self._save_dir is None:
197            return
198        from PIL import Image
199
200        meta = getattr(sample, "metadata", None) or {}
201        scene = meta.get("scene_id")
202        phase = meta.get("phase_idx")
203        frame = meta.get("frame")
204        if not (scene and frame is not None and phase is not None):
205            try:
206                scene, phase, frame = str(sample.id).split("__", 2)
207            except ValueError:
208                scene, phase, frame = "unknown", "0", str(sample.id)
209        out = self._save_dir / str(scene) / str(phase) / f"{frame}.png"
210        out.parent.mkdir(parents=True, exist_ok=True)
211        # 16-bit PNG preserves up to 65 K instance ids; sufficient for SAM2.
212        Image.fromarray(np.asarray(model_output, dtype=np.uint16), mode="I;16").save(out)
213
214
215class BatchedRelativePoseBenchmarkModel(BatchedTaskBenchmarkModel):
216    """Batched wrapper for paired-frame relative-pose models.
217
218    Adapter contract: ``adapter(list_of_dicts) -> list_of_pose_dicts``,
219    where each input dict is ``{"rgb_a": HxWx3 uint8, "rgb_b": HxWx3 uint8}``
220    and each output is ``{"rotation": 3x3, "translation": 3}`` (or a tuple).
221    """
222
223    task = TaskType.RELATIVE_CAMERA_POSE
224
225    def extract_inputs(self, sample: Sample) -> dict:
226        meta = getattr(sample, "metadata", None) or {}
227        rgb_a = np.asarray(sample.rgb, dtype=np.uint8)
228        rgb_b = meta.get("rgb_b")
229        if rgb_b is None:
230            from ..exceptions import AdapterError
231
232            raise AdapterError(
233                "Relative-pose adapter requires `rgb_b` in sample.metadata; "
234                "check the manifest writes `rgb_b` so the loader stashes it.",
235            )
236        return {"rgb_a": rgb_a, "rgb_b": np.asarray(rgb_b, dtype=np.uint8)}
237
238    def wrap_output(self, model_output: Any, sample: Sample) -> Any:
239        from ..api import RelativePosePrediction
240
241        if isinstance(model_output, dict):
242            rot = model_output["rotation"]
243            trans = model_output["translation"]
244        else:
245            rot, trans = model_output
246        return RelativePosePrediction(
247            rotation=np.asarray(rot, dtype=np.float64),
248            translation=np.asarray(trans, dtype=np.float64),
249        )
250
251    #: CSV columns. ``sample_id`` is the canonical identity; both phases are
252    #: retained so cross-phase pairs cannot collide with intra-phase pairs.
253    _LOG_COLUMNS: tuple[str, ...] = (
254        "sample_id",
255        "scene_id",
256        "phase",
257        "phase_b",
258        "frame_a",
259        "frame_b",
260        "R00",
261        "R01",
262        "R02",
263        "R10",
264        "R11",
265        "R12",
266        "R20",
267        "R21",
268        "R22",
269        "tx",
270        "ty",
271        "tz",
272    )
273
274    def maybe_save(self, sample: Sample, model_output: Any) -> None:
275        """Append one row per pair to ``<save_dir>/predictions.csv``.
276
277        Rationale: paired-pose runs produce ~3 K rows per (model, split);
278        a single CSV is friendlier for Box and analytics than 3 K
279        per-pair ``.npz`` files. Columns: scene_id, phase, frame_a,
280        frame_b, the 9 rotation entries (row-major) and 3 translation
281        components. Header is written exactly once per file (on the
282        first append to an empty / non-existent log). Idempotent on
283        ``sample.id`` within a single process so
284        the runner's deployment-readiness double pass (warmup + measure)
285        doesn't duplicate rows.
286        """
287        if self._save_dir is None:
288            return
289        import csv
290
291        meta = getattr(sample, "metadata", None) or {}
292        scene = meta.get("scene_id") or "unknown"
293        phase = meta.get("phase_idx") if meta.get("phase_idx") is not None else "0"
294        phase_b = meta.get("phase_idx_b")
295        if phase_b is None:
296            phase_b = phase
297        frame_a = meta.get("frame") or str(sample.id)
298        frame_b = meta.get("frame_b") or "?"
299        sample_id = str(sample.id)
300
301        seen = self.__dict__.setdefault("_pose_csv_seen", set())
302        key = sample_id
303        if key in seen:
304            return
305        seen.add(key)
306
307        if isinstance(model_output, dict):
308            rot = np.asarray(model_output["rotation"], dtype=np.float64)
309            trans = np.asarray(model_output["translation"], dtype=np.float64)
310        else:
311            rot, trans = model_output
312            rot = np.asarray(rot, dtype=np.float64)
313            trans = np.asarray(trans, dtype=np.float64)
314
315        # Tolerate rotation given as a flat 9-vector or a quaternion.
316        rot = rot.reshape(-1)
317        if rot.size == 4:
318            # Quaternion (x, y, z, w) → 3×3 rotation. We don't depend on
319            # scipy here; build it inline from the unit-quaternion formula.
320            x, y, z, w = rot.tolist()
321            rot = np.asarray(
322                [
323                    1 - 2 * (y * y + z * z),
324                    2 * (x * y - z * w),
325                    2 * (x * z + y * w),
326                    2 * (x * y + z * w),
327                    1 - 2 * (x * x + z * z),
328                    2 * (y * z - x * w),
329                    2 * (x * z - y * w),
330                    2 * (y * z + x * w),
331                    1 - 2 * (x * x + y * y),
332                ],
333                dtype=np.float64,
334            )
335        elif rot.size != 9:
336            from ..exceptions import AdapterError
337
338            raise AdapterError(
339                f"unexpected rotation shape: got {rot.size} entries; "
340                "expected 9 (3×3 matrix) or 4 (xyzw quaternion).",
341            )
342        trans = trans.reshape(-1)
343        if trans.size != 3:
344            from ..exceptions import AdapterError
345
346            raise AdapterError(
347                f"unexpected translation shape: got {trans.size} entries; expected 3.",
348            )
349
350        log_path = self._save_dir / "predictions.csv"
351        log_path.parent.mkdir(parents=True, exist_ok=True)
352        write_header = not log_path.exists() or log_path.stat().st_size == 0
353        with log_path.open("a", newline="") as f:
354            writer = csv.writer(f)
355            if write_header:
356                writer.writerow(self._LOG_COLUMNS)
357            writer.writerow(
358                [
359                    sample_id,
360                    scene,
361                    phase,
362                    phase_b,
363                    frame_a,
364                    frame_b,
365                    *(f"{v:.10g}" for v in rot.tolist()),
366                    *(f"{v:.10g}" for v in trans.tolist()),
367                ]
368            )