rpx_benchmark.adapters.base

   1"""Core types for the RPX adapter framework.
   2
   3::
   4
   5    Sample ─► InputAdapter.prepare ─► PreparedInput(payload, context)
   6                                           │
   7                                           ▼
   8                                       model(payload)
   9                                           │
  10                                           ▼
  11    Sample, context, model_output ─► OutputAdapter.finalize ─► Prediction
  12
  13Users extending RPX for their own model only need to supply the *model*
  14and pick a matching pair of adapters. The adapters ship with the
  15library for common families (HuggingFace transformers, UniDepth,
  16Metric3D, raw numpy callables).
  17"""
  18
  19from __future__ import annotations
  20
  21from dataclasses import dataclass, field
  22from typing import Any, Callable, Dict, Optional, Protocol, Sequence, runtime_checkable
  23
  24import numpy as np
  25
  26from ..api import (
  27    BenchmarkModel,
  28    DepthPrediction,
  29    DetectionPrediction,
  30    KeypointCorrespondencePrediction,
  31    RelativePosePrediction,
  32    Sample,
  33    SegmentationPrediction,
  34    SparseDepthPrediction,
  35    TaskType,
  36    Tracklet,
  37    TrackletPrediction,
  38    VideoDepthPrediction,
  39    VideoSample,
  40    VisualGroundingPrediction,
  41)
  42from ..exceptions import AdapterError
  43
  44# --------------------------------------------------------------------------- #
  45# Payload container
  46# --------------------------------------------------------------------------- #
  47
  48
  49@dataclass
  50class PreparedInput:
  51    """Everything a model needs for one sample, plus context for post-processing.
  52
  53    ``payload`` is whatever the model's forward call accepts. If it is a
  54    ``dict``, the default invoker calls ``model(**payload)``; otherwise
  55    ``model(payload)``.
  56
  57    ``context`` is a free-form dict the output adapter receives back. Use
  58    it to stash things like target image size, original intrinsics, or
  59    any preprocessing metadata the postprocessing step needs.
  60    """
  61
  62    payload: Any
  63    context: Dict[str, Any] = field(default_factory=dict)
  64
  65
  66# --------------------------------------------------------------------------- #
  67# Protocol types
  68# --------------------------------------------------------------------------- #
  69
  70
  71@runtime_checkable
  72class InputAdapter(Protocol):
  73    """Sample → model-ready payload."""
  74
  75    def setup(self) -> None:  # pragma: no cover - optional hook
  76        """Optional one-time setup (e.g., build a processor on first use)."""
  77
  78    def prepare(self, sample: Sample) -> PreparedInput: ...
  79
  80
  81@runtime_checkable
  82class OutputAdapter(Protocol):
  83    """Model output → RPX prediction object."""
  84
  85    def setup(self) -> None:  # pragma: no cover - optional hook
  86        """Optional one-time setup."""
  87
  88    def finalize(
  89        self,
  90        model_output: Any,
  91        context: Dict[str, Any],
  92        sample: Sample,
  93    ) -> Any: ...
  94
  95
  96ModelInvoker = Callable[[Any, Any], Any]
  97"""Given (model, payload) return the raw model output.
  98
  99The default implementation handles the two common cases: dict payloads
 100become ``model(**payload)``, anything else becomes ``model(payload)``.
 101Override when the model needs non-standard invocation (e.g., a method
 102other than ``__call__``).
 103"""
 104
 105
 106def default_invoker(model: Any, payload: Any) -> Any:
 107    """Call ``model`` with ``payload``. Uses torch no_grad if available."""
 108    try:
 109        import torch  # noqa: F401
 110    except ImportError:
 111        return _dispatch(model, payload)
 112
 113    import torch
 114
 115    with torch.no_grad():
 116        return _dispatch(model, payload)
 117
 118
 119def _dispatch(model: Any, payload: Any) -> Any:
 120    if isinstance(payload, dict):
 121        return model(**payload)
 122    return model(payload)
 123
 124
 125# --------------------------------------------------------------------------- #
 126# BenchmarkableModel
 127# --------------------------------------------------------------------------- #
 128
 129
 130@dataclass
 131class BenchmarkableModel(BenchmarkModel):
 132    """Compose an input adapter, a model, and an output adapter.
 133
 134    This is the canonical way to plug a model into the RPX benchmark
 135    harness. The ``BenchmarkRunner`` only ever sees the
 136    :class:`BenchmarkModel` contract (``setup``, ``predict``); all the
 137    model-family-specific logic lives in the adapters so the harness
 138    stays task-agnostic.
 139
 140    Example — wrap a plain numpy depth callable::
 141
 142        import numpy as np
 143        import rpx_benchmark as rpx
 144
 145        def my_depth(rgb: np.ndarray) -> np.ndarray:
 146            return ...           # H x W float32, metres
 147
 148        bm = rpx.make_numpy_depth_model(my_depth, name="my_model")
 149        # Hand `bm` to a task runner or the BenchmarkRunner directly.
 150
 151    Example — wrap a plain numpy callable::
 152
 153        def my_depth(rgb: np.ndarray) -> np.ndarray:
 154            return some_depth_in_metres
 155
 156        bm = make_numpy_depth_model(my_depth)
 157    """
 158
 159    task: TaskType
 160    input_adapter: InputAdapter
 161    model: Any
 162    output_adapter: OutputAdapter
 163    invoker: ModelInvoker = field(default=default_invoker)
 164    name: str = "benchmarkable_model"
 165    setup_hook: Optional[Callable[[], None]] = None
 166
 167    def setup(self) -> None:
 168        if hasattr(self.input_adapter, "setup"):
 169            self.input_adapter.setup()
 170        if hasattr(self.output_adapter, "setup"):
 171            self.output_adapter.setup()
 172        if self.setup_hook is not None:
 173            self.setup_hook()
 174
 175    def predict(self, batch: Sequence[Sample]) -> Sequence[Any]:
 176        preds = []
 177        for sample in batch:
 178            prepared = self.input_adapter.prepare(sample)
 179            model_out = self.invoker(self.model, prepared.payload)
 180            pred = self.output_adapter.finalize(model_out, prepared.context, sample)
 181            preds.append(pred)
 182        return preds
 183
 184
 185# --------------------------------------------------------------------------- #
 186# Convenience factory: numpy-in / numpy-out depth model
 187# --------------------------------------------------------------------------- #
 188
 189
 190class _NumpyDepthInput:
 191    """Hand the model the raw ``(H, W, 3) uint8`` array untouched."""
 192
 193    def prepare(self, sample: Sample) -> PreparedInput:
 194        rgb = np.asarray(sample.rgb, dtype=np.uint8)
 195        return PreparedInput(payload=rgb, context={"target_hw": rgb.shape[:2]})
 196
 197
 198class _NumpyDepthOutput:
 199    """Wrap a raw ``(H_pred, W_pred) float`` array into ``DepthPrediction``.
 200
 201    If the model returns a shape different from ``target_hw``, the result
 202    is bilinearly resized (via PIL) so ``BenchmarkRunner`` can compare
 203    pixel-for-pixel against the GT.
 204    """
 205
 206    def finalize(
 207        self,
 208        model_output: Any,
 209        context: Dict[str, Any],
 210        sample: Sample,
 211    ) -> DepthPrediction:
 212        depth = np.asarray(model_output, dtype=np.float32)
 213        if depth.ndim != 2:
 214            raise AdapterError(
 215                f"numpy depth model must return a 2-D array; got shape {depth.shape}",
 216                hint="Your depth callable must return a (H, W) float array. "
 217                "A common mistake is returning (1, H, W) or (H, W, 1).",
 218            )
 219        target_hw = context.get("target_hw", depth.shape)
 220        if depth.shape != target_hw:
 221            from PIL import Image
 222
 223            pil = Image.fromarray(depth, mode="F")
 224            pil = pil.resize((target_hw[1], target_hw[0]), Image.BILINEAR)
 225            depth = np.asarray(pil, dtype=np.float32)
 226        return DepthPrediction(depth_map=depth)
 227
 228
 229# --------------------------------------------------------------------------- #
 230# Convenience factory: numpy-in / numpy-out video-depth model
 231# --------------------------------------------------------------------------- #
 232
 233
 234class _NumpyVideoDepthInput:
 235    """Hand the model the raw ``(T, H, W, 3) uint8`` clip untouched."""
 236
 237    def prepare(self, sample: VideoSample) -> PreparedInput:
 238        rgb_seq = np.asarray(sample.rgb_seq, dtype=np.uint8)
 239        return PreparedInput(
 240            payload=rgb_seq,
 241            context={"target_thw": rgb_seq.shape[:3]},
 242        )
 243
 244
 245class _NumpyVideoDepthOutput:
 246    """Wrap a raw ``(T, H_pred, W_pred) float`` array into
 247    :class:`VideoDepthPrediction`.
 248
 249    If the model returns a shape different from ``target_thw``, the
 250    per-frame depth maps are bilinearly resized (via PIL) so the
 251    runner can compare pixel-for-pixel against the GT clip. T must
 252    match — a model that drops frames is a contract violation, not
 253    something this output adapter is allowed to paper over.
 254    """
 255
 256    def finalize(
 257        self,
 258        model_output: Any,
 259        context: Dict[str, Any],
 260        sample: VideoSample,
 261    ) -> VideoDepthPrediction:
 262        depth_seq = np.asarray(model_output, dtype=np.float32)
 263        if depth_seq.ndim == 4 and depth_seq.shape[1] == 1:
 264            # (T, 1, H, W) → (T, H, W); a common model output shape.
 265            depth_seq = depth_seq.squeeze(1)
 266        if depth_seq.ndim != 3:
 267            raise AdapterError(
 268                f"numpy video-depth model must return a 3-D (T, H, W) "
 269                f"array; got shape {depth_seq.shape}",
 270                hint="Your video-depth callable must return (T, H, W) "
 271                "float32. A common mistake is returning (T, 1, H, W) "
 272                "without squeezing the channel dim.",
 273            )
 274        target_thw = context.get("target_thw", depth_seq.shape)
 275        target_t, target_h, target_w = target_thw
 276        if depth_seq.shape[0] != target_t:
 277            raise AdapterError(
 278                f"numpy video-depth model returned {depth_seq.shape[0]} "
 279                f"frames; clip has {target_t}. The model must emit one "
 280                "depth map per input frame — runner / metric calculators "
 281                "assume aligned T.",
 282            )
 283        if depth_seq.shape[1:] != (target_h, target_w):
 284            from PIL import Image
 285
 286            resized = np.empty(
 287                (target_t, target_h, target_w), dtype=np.float32
 288            )
 289            for t in range(target_t):
 290                pil = Image.fromarray(depth_seq[t], mode="F")
 291                pil = pil.resize((target_w, target_h), Image.BILINEAR)
 292                resized[t] = np.asarray(pil, dtype=np.float32)
 293            depth_seq = resized
 294        return VideoDepthPrediction(depth_map_seq=depth_seq)
 295
 296
 297def make_numpy_video_depth_model(
 298    fn: Callable[[np.ndarray], np.ndarray],
 299    *,
 300    name: str = "numpy_video_depth_model",
 301    depth_output_kind: str = "metric",
 302) -> BenchmarkableModel:
 303    """Wrap a plain numpy video-depth callable as a
 304    :class:`BenchmarkableModel`.
 305
 306    The callable must accept a ``(T, H, W, 3) uint8`` RGB clip and
 307    return a ``(T, H', W') float`` per-frame depth sequence (metres
 308    if metric; up-to-scale if relative — the runner applies per-clip
 309    ``(s, t)`` alignment for relative models). If ``(H', W') !=
 310    (H, W)``, the output is bilinearly resized per frame to match.
 311    ``T`` must match the input clip length.
 312
 313    Parameters
 314    ----------
 315    fn : callable
 316        Signature: ``fn(rgb_seq_uint8) -> depth_seq_float``.
 317    name : str
 318        Display name used in logs and reports.
 319    depth_output_kind : str
 320        ``"metric"`` (default) — model emits metres; runner skips
 321        alignment. ``"relative"`` — runner applies per-clip
 322        ``(s, t)`` scale-and-shift alignment before metrics.
 323
 324    Examples
 325    --------
 326    >>> import numpy as np
 327    >>> import rpx_benchmark as rpx
 328    >>> def my_clip_depth(rgb_seq):
 329    ...     T, H, W, _ = rgb_seq.shape
 330    ...     return np.full((T, H, W), 2.0, dtype=np.float32)
 331    >>> bm = rpx.make_numpy_video_depth_model(my_clip_depth, name="mine")
 332    >>> bm.task is rpx.TaskType.VIDEO_DEPTH
 333    True
 334    """
 335    model = BenchmarkableModel(
 336        task=TaskType.VIDEO_DEPTH,
 337        input_adapter=_NumpyVideoDepthInput(),
 338        model=fn,
 339        output_adapter=_NumpyVideoDepthOutput(),
 340        invoker=lambda model, payload: model(payload),
 341        name=name,
 342    )
 343    # Propagate depth_output_kind to the runner's alignment dispatch.
 344    model.depth_output_kind = depth_output_kind
 345    return model
 346
 347
 348def make_numpy_depth_model(
 349    fn: Callable[[np.ndarray], np.ndarray],
 350    *,
 351    name: str = "numpy_depth_model",
 352    depth_output_kind: str = "metric",
 353    native_alignment: str | None = None,
 354) -> BenchmarkableModel:
 355    """Wrap a plain numpy depth callable as a :class:`BenchmarkableModel`.
 356
 357    The callable must accept a ``(H, W, 3) uint8`` RGB image and return a
 358    ``(H', W') float`` depth map. If ``depth_output_kind="metric"`` the
 359    values are metres. Relative outputs remain raw until the runner fits
 360    one pooled alignment per ``(scene, phase)`` cell. If ``(H', W') !=
 361    (H, W)``, the output is bilinearly resized to match the ground truth.
 362
 363    Parameters
 364    ----------
 365    fn : callable
 366        The depth function. Signature: ``fn(rgb_uint8) -> depth_float``.
 367    name : str
 368        Display name used in logs and reports.
 369    depth_output_kind : str
 370        ``"metric"`` (default) or ``"relative"``.
 371    native_alignment : str, optional
 372        Alignment used for relative output. Defaults to ``"ls_affine"``
 373        for relative models and ``"none"`` for metric models.
 374
 375    Examples
 376    --------
 377    >>> import numpy as np
 378    >>> import rpx_benchmark as rpx
 379    >>> def my_depth(rgb):
 380    ...     return np.full(rgb.shape[:2], 2.0, dtype=np.float32)
 381    >>> bm = rpx.make_numpy_depth_model(my_depth, name="mine")
 382    >>> bm.task is rpx.TaskType.MONOCULAR_DEPTH
 383    True
 384    """
 385    if depth_output_kind not in {"metric", "relative"}:
 386        raise AdapterError(
 387            f"depth_output_kind must be 'metric' or 'relative', got {depth_output_kind!r}",
 388        )
 389    model = BenchmarkableModel(
 390        task=TaskType.MONOCULAR_DEPTH,
 391        input_adapter=_NumpyDepthInput(),
 392        model=fn,
 393        output_adapter=_NumpyDepthOutput(),
 394        invoker=lambda model, payload: model(payload),
 395        name=name,
 396    )
 397    model.depth_output_kind = depth_output_kind
 398    model.native_alignment = native_alignment or (
 399        "ls_affine" if depth_output_kind == "relative" else "none"
 400    )
 401    return model
 402
 403
 404# --------------------------------------------------------------------------- #
 405# Convenience factory: numpy-in / int-mask-out segmentation model
 406# --------------------------------------------------------------------------- #
 407
 408
 409class _NumpyMaskInput:
 410    """Hand the model the raw ``(H, W, 3) uint8`` array untouched."""
 411
 412    def prepare(self, sample: Sample) -> PreparedInput:
 413        rgb = np.asarray(sample.rgb, dtype=np.uint8)
 414        return PreparedInput(payload=rgb, context={"target_hw": rgb.shape[:2]})
 415
 416
 417class _NumpyMaskOutput:
 418    """Wrap a raw ``(H_pred, W_pred) int`` array into ``SegmentationPrediction``.
 419
 420    If the model returns a shape different from ``target_hw``, the result
 421    is nearest-neighbour resized (masks must keep integer IDs — bilinear
 422    would introduce invalid fractional values).
 423    """
 424
 425    def finalize(
 426        self,
 427        model_output: Any,
 428        context: Dict[str, Any],
 429        sample: Sample,
 430    ) -> SegmentationPrediction:
 431        mask = np.asarray(model_output)
 432        if mask.ndim != 2:
 433            raise AdapterError(
 434                f"numpy mask model must return a 2-D array; got shape {mask.shape}",
 435                hint="Your mask callable must return a (H, W) int array. "
 436                "A common mistake is returning (H, W, C) per-class logits.",
 437            )
 438        mask = mask.astype(np.int32, copy=False)
 439        target_hw = context.get("target_hw", mask.shape)
 440        if mask.shape != target_hw:
 441            from PIL import Image
 442
 443            pil = Image.fromarray(mask.astype(np.int32), mode="I")
 444            pil = pil.resize((target_hw[1], target_hw[0]), Image.NEAREST)
 445            mask = np.asarray(pil, dtype=np.int32)
 446        return SegmentationPrediction(mask=mask)
 447
 448
 449def make_numpy_mask_model(
 450    fn: Callable[[np.ndarray], np.ndarray],
 451    *,
 452    name: str = "numpy_mask_model",
 453) -> BenchmarkableModel:
 454    """Wrap a plain numpy instance-mask callable as a :class:`BenchmarkableModel`.
 455
 456    The callable must accept a ``(H, W, 3) uint8`` RGB image and return
 457    a ``(H', W') int`` instance mask where pixel values are instance
 458    IDs (``0`` is conventionally background). If ``(H', W') != (H, W)``
 459    the output is nearest-neighbour resized to match the GT mask so
 460    integer IDs are preserved.
 461
 462    Parameters
 463    ----------
 464    fn : callable
 465        ``fn(rgb_uint8) -> mask_int``.
 466    name : str
 467
 468    Examples
 469    --------
 470    >>> import numpy as np
 471    >>> import rpx_benchmark as rpx
 472    >>> def my_seg(rgb):
 473    ...     mask = np.zeros(rgb.shape[:2], dtype=np.int32)
 474    ...     mask[rgb.sum(-1) > 384] = 1  # trivial brightness threshold
 475    ...     return mask
 476    >>> bm = rpx.make_numpy_mask_model(my_seg, name="mine")
 477    >>> bm.task is rpx.TaskType.OBJECT_SEGMENTATION
 478    True
 479    """
 480    return BenchmarkableModel(
 481        task=TaskType.OBJECT_SEGMENTATION,
 482        input_adapter=_NumpyMaskInput(),
 483        model=fn,
 484        output_adapter=_NumpyMaskOutput(),
 485        invoker=lambda model, payload: model(payload),
 486        name=name,
 487    )
 488
 489
 490# --------------------------------------------------------------------------- #
 491# Shared helpers for numpy fast paths
 492# --------------------------------------------------------------------------- #
 493
 494
 495def _passthrough_invoker(model: Any, payload: Any) -> Any:
 496    """Invoker for numpy fast paths — just call the function on its payload."""
 497    return model(payload)
 498
 499
 500class _NumpyRgbInput:
 501    """Every numpy task's input adapter: just hand the RGB ndarray through.
 502
 503    Extra tokens (second RGB, target pose, ...) travel in
 504    ``sample.metadata`` and are picked out by the per-task output
 505    adapter, so a single input adapter serves every task.
 506    """
 507
 508    def prepare(self, sample: Sample) -> PreparedInput:
 509        rgb = np.asarray(sample.rgb, dtype=np.uint8)
 510        context = {"target_hw": rgb.shape[:2], "sample": sample}
 511        return PreparedInput(payload=rgb, context=context)
 512
 513
 514def _coerce_dict_result(
 515    result: Any,
 516    *,
 517    required_keys: Sequence[str],
 518    task_name: str,
 519) -> Dict[str, Any]:
 520    """Normalise a function's return value into a dict with the expected keys.
 521
 522    Accepts either a dict or an ``(a, b, ...)`` tuple aligned with
 523    ``required_keys``. Raises :class:`AdapterError` when neither
 524    shape matches.
 525    """
 526    if isinstance(result, dict):
 527        missing = [k for k in required_keys if k not in result]
 528        if missing:
 529            raise AdapterError(
 530                f"{task_name} model result is missing required keys: {missing}",
 531                hint=f"Return a dict with keys {list(required_keys)}.",
 532            )
 533        return {k: result[k] for k in required_keys}
 534    if isinstance(result, tuple) and len(result) == len(required_keys):
 535        return dict(zip(required_keys, result, strict=False))
 536    raise AdapterError(
 537        f"{task_name} model must return a dict with keys "
 538        f"{list(required_keys)} or a {len(required_keys)}-tuple in that order; "
 539        f"got {type(result).__name__}",
 540    )
 541
 542
 543# --------------------------------------------------------------------------- #
 544# Object detection — fn(rgb) -> {boxes, scores, labels}
 545# --------------------------------------------------------------------------- #
 546
 547
 548class _NumpyDetectionOutput:
 549    """Wrap a detection callable's output into a :class:`DetectionPrediction`."""
 550
 551    def finalize(
 552        self,
 553        model_output: Any,
 554        context: Dict[str, Any],
 555        sample: Sample,
 556    ) -> DetectionPrediction:
 557        fields = _coerce_dict_result(
 558            model_output,
 559            required_keys=("boxes", "scores", "labels"),
 560            task_name="detection",
 561        )
 562        boxes = np.asarray(fields["boxes"], dtype=np.float32).reshape(-1, 4)
 563        scores = np.asarray(fields["scores"], dtype=np.float32).reshape(-1)
 564        labels = list(fields["labels"])
 565        if not (len(boxes) == len(scores) == len(labels)):
 566            raise AdapterError(
 567                f"detection model returned mismatched lengths: "
 568                f"boxes={len(boxes)} scores={len(scores)} labels={len(labels)}",
 569            )
 570        return DetectionPrediction(boxes=boxes, scores=scores, labels=labels)
 571
 572
 573def make_numpy_detection_model(
 574    fn: Callable[[np.ndarray], Any],
 575    *,
 576    name: str = "numpy_detection_model",
 577    task: TaskType = TaskType.OBJECT_DETECTION,
 578) -> BenchmarkableModel:
 579    """Wrap a plain numpy detection callable as a :class:`BenchmarkableModel`.
 580
 581    The callable must accept a ``(H, W, 3) uint8`` RGB image and
 582    return either a dict with keys ``"boxes"`` / ``"scores"`` /
 583    ``"labels"`` or a ``(boxes, scores, labels)`` tuple in that
 584    order. Boxes are pixel coordinates in ``(x1, y1, x2, y2)``
 585    format, scores are floats in ``[0, 1]``, labels are strings.
 586
 587    Parameters
 588    ----------
 589    fn : callable
 590        ``fn(rgb_uint8) -> dict | tuple``.
 591    name : str
 592        Display name for reports.
 593    task : TaskType
 594        Use :attr:`TaskType.OBJECT_DETECTION` for closed-vocabulary
 595        detection or :attr:`TaskType.OPEN_VOCAB_DETECTION` for
 596        open-vocab (the Prediction contract is the same).
 597
 598    Examples
 599    --------
 600    >>> import numpy as np
 601    >>> import rpx_benchmark as rpx
 602    >>> def my_det(rgb):
 603    ...     return {
 604    ...         "boxes": np.array([[10, 10, 30, 30]], dtype=np.float32),
 605    ...         "scores": np.array([0.9], dtype=np.float32),
 606    ...         "labels": ["cup"],
 607    ...     }
 608    >>> bm = rpx.make_numpy_detection_model(my_det)
 609    >>> bm.task is rpx.TaskType.OBJECT_DETECTION
 610    True
 611    """
 612    return BenchmarkableModel(
 613        task=task,
 614        input_adapter=_NumpyRgbInput(),
 615        model=fn,
 616        output_adapter=_NumpyDetectionOutput(),
 617        invoker=_passthrough_invoker,
 618        name=name,
 619    )
 620
 621
 622# --------------------------------------------------------------------------- #
 623# Visual grounding — fn(rgb, text) -> {boxes, scores}
 624# --------------------------------------------------------------------------- #
 625
 626
 627class _NumpyGroundingInput:
 628    """Grounding input: raw RGB + the referring expression from the GT."""
 629
 630    def prepare(self, sample: Sample) -> PreparedInput:
 631        rgb = np.asarray(sample.rgb, dtype=np.uint8)
 632        text = getattr(sample.ground_truth, "text", "")
 633        return PreparedInput(
 634            payload={"rgb": rgb, "text": text},
 635            context={"target_hw": rgb.shape[:2]},
 636        )
 637
 638
 639class _NumpyGroundingOutput:
 640    def finalize(
 641        self,
 642        model_output: Any,
 643        context: Dict[str, Any],
 644        sample: Sample,
 645    ) -> VisualGroundingPrediction:
 646        fields = _coerce_dict_result(
 647            model_output,
 648            required_keys=("boxes", "scores"),
 649            task_name="grounding",
 650        )
 651        boxes = np.asarray(fields["boxes"], dtype=np.float32).reshape(-1, 4)
 652        scores = np.asarray(fields["scores"], dtype=np.float32).reshape(-1)
 653        labels = fields.get("labels") if isinstance(fields, dict) else None
 654        return VisualGroundingPrediction(
 655            boxes=boxes,
 656            scores=scores,
 657            labels=labels,
 658        )
 659
 660
 661def _grounding_invoker(model: Any, payload: Dict[str, Any]) -> Any:
 662    """Grounding callable signature: ``fn(rgb, text)``."""
 663    return model(payload["rgb"], payload["text"])
 664
 665
 666def make_numpy_grounding_model(
 667    fn: Callable[[np.ndarray, str], Any],
 668    *,
 669    name: str = "numpy_grounding_model",
 670) -> BenchmarkableModel:
 671    """Wrap a visual-grounding callable as a :class:`BenchmarkableModel`.
 672
 673    The callable takes ``(rgb_uint8, text)`` and returns either a
 674    dict with keys ``"boxes"`` / ``"scores"`` or a tuple
 675    ``(boxes, scores)``. Boxes are ``(x1, y1, x2, y2)`` pixel
 676    coordinates; scores are floats. The referring expression
 677    ``text`` is plucked from ``sample.ground_truth.text`` by the
 678    adapter so the callable never sees the GT boxes.
 679
 680    Examples
 681    --------
 682    >>> import numpy as np
 683    >>> import rpx_benchmark as rpx
 684    >>> def my_ground(rgb, text):
 685    ...     return (
 686    ...         np.array([[10, 10, 30, 30]], dtype=np.float32),
 687    ...         np.array([0.8], dtype=np.float32),
 688    ...     )
 689    >>> bm = rpx.make_numpy_grounding_model(my_ground)
 690    """
 691    return BenchmarkableModel(
 692        task=TaskType.VISUAL_GROUNDING,
 693        input_adapter=_NumpyGroundingInput(),
 694        model=fn,
 695        output_adapter=_NumpyGroundingOutput(),
 696        invoker=_grounding_invoker,
 697        name=name,
 698    )
 699
 700
 701# --------------------------------------------------------------------------- #
 702# Relative camera pose — fn(rgb_a, rgb_b) -> {rotation, translation}
 703# --------------------------------------------------------------------------- #
 704
 705
 706class _NumpyPoseInput:
 707    """Pose input: extract rgb_a + rgb_b from the sample metadata side-channel."""
 708
 709    def prepare(self, sample: Sample) -> PreparedInput:
 710        rgb_a = np.asarray(sample.rgb, dtype=np.uint8)
 711        rgb_b = None
 712        if sample.metadata is not None:
 713            rgb_b = sample.metadata.get("rgb_b")
 714        if rgb_b is None:
 715            raise AdapterError(
 716                "Relative-pose numpy adapter requires a second RGB frame "
 717                "(`rgb_b`) loaded by the dataset.",
 718                hint="Check that your manifest entries include `rgb_b` "
 719                "alongside `rgb` so the loader stashes it in sample.metadata.",
 720            )
 721        return PreparedInput(
 722            payload={"rgb_a": rgb_a, "rgb_b": np.asarray(rgb_b, dtype=np.uint8)},
 723            context={"target_hw": rgb_a.shape[:2]},
 724        )
 725
 726
 727class _NumpyPoseOutput:
 728    def finalize(
 729        self,
 730        model_output: Any,
 731        context: Dict[str, Any],
 732        sample: Sample,
 733    ) -> RelativePosePrediction:
 734        fields = _coerce_dict_result(
 735            model_output,
 736            required_keys=("rotation", "translation"),
 737            task_name="relative_pose",
 738        )
 739        return RelativePosePrediction(
 740            rotation=np.asarray(fields["rotation"], dtype=np.float64),
 741            translation=np.asarray(fields["translation"], dtype=np.float64),
 742        )
 743
 744
 745def _pose_invoker(model: Any, payload: Dict[str, Any]) -> Any:
 746    return model(payload["rgb_a"], payload["rgb_b"])
 747
 748
 749def make_numpy_pose_model(
 750    fn: Callable[[np.ndarray, np.ndarray], Any],
 751    *,
 752    name: str = "numpy_pose_model",
 753) -> BenchmarkableModel:
 754    """Wrap a relative-camera-pose callable as a :class:`BenchmarkableModel`.
 755
 756    The callable takes ``(rgb_a, rgb_b)`` and returns either a dict
 757    with keys ``"rotation"`` (3×3 rotation matrix or 4-element
 758    quaternion) and ``"translation"`` (3-vector, metres) or a
 759    ``(rotation, translation)`` tuple.
 760
 761    Examples
 762    --------
 763    >>> import numpy as np
 764    >>> import rpx_benchmark as rpx
 765    >>> def my_pose(rgb_a, rgb_b):
 766    ...     return {"rotation": np.eye(3), "translation": np.zeros(3)}
 767    >>> bm = rpx.make_numpy_pose_model(my_pose)
 768    """
 769    return BenchmarkableModel(
 770        task=TaskType.RELATIVE_CAMERA_POSE,
 771        input_adapter=_NumpyPoseInput(),
 772        model=fn,
 773        output_adapter=_NumpyPoseOutput(),
 774        invoker=_pose_invoker,
 775        name=name,
 776    )
 777
 778
 779# --------------------------------------------------------------------------- #
 780# Sparse depth — fn(rgb, coords) -> depths
 781# --------------------------------------------------------------------------- #
 782
 783
 784class _NumpySparseDepthInput:
 785    """Sparse-depth input: rgb + the GT sparse pixel coordinates."""
 786
 787    def prepare(self, sample: Sample) -> PreparedInput:
 788        rgb = np.asarray(sample.rgb, dtype=np.uint8)
 789        gt = sample.ground_truth
 790        coords = getattr(gt, "coordinates", None)
 791        if coords is None:
 792            raise AdapterError(
 793                "Sparse-depth ground truth missing `coordinates`.",
 794            )
 795        return PreparedInput(
 796            payload={"rgb": rgb, "coordinates": np.asarray(coords, dtype=np.float32)},
 797            context={"target_hw": rgb.shape[:2]},
 798        )
 799
 800
 801class _NumpySparseDepthOutput:
 802    def finalize(
 803        self,
 804        model_output: Any,
 805        context: Dict[str, Any],
 806        sample: Sample,
 807    ) -> SparseDepthPrediction:
 808        # The callable is responsible for returning depths at the
 809        # sparse coords it was given; we pass its coords back as-is.
 810        coords = np.asarray(sample.ground_truth.coordinates, dtype=np.float32)
 811        depths = np.asarray(model_output, dtype=np.float32).reshape(-1)
 812        if len(depths) != len(coords):
 813            raise AdapterError(
 814                f"sparse-depth model returned {len(depths)} depths for "
 815                f"{len(coords)} queried coordinates.",
 816            )
 817        return SparseDepthPrediction(coordinates=coords, depths=depths)
 818
 819
 820def _sparse_depth_invoker(model: Any, payload: Dict[str, Any]) -> Any:
 821    return model(payload["rgb"], payload["coordinates"])
 822
 823
 824def make_numpy_sparse_depth_model(
 825    fn: Callable[[np.ndarray, np.ndarray], np.ndarray],
 826    *,
 827    name: str = "numpy_sparse_depth_model",
 828) -> BenchmarkableModel:
 829    """Wrap a sparse-depth callable as a :class:`BenchmarkableModel`.
 830
 831    The callable takes ``(rgb_uint8, coords)`` where ``coords`` is a
 832    ``(N, 2)`` float32 array of pixel coordinates and returns an
 833    ``(N,)`` float32 array of depths in metres at those exact
 834    coordinates.
 835
 836    Examples
 837    --------
 838    >>> import numpy as np
 839    >>> import rpx_benchmark as rpx
 840    >>> def my_sd(rgb, coords):
 841    ...     return np.full(len(coords), 2.0, dtype=np.float32)
 842    >>> bm = rpx.make_numpy_sparse_depth_model(my_sd)
 843    """
 844    return BenchmarkableModel(
 845        task=TaskType.SPARSE_DEPTH,
 846        input_adapter=_NumpySparseDepthInput(),
 847        model=fn,
 848        output_adapter=_NumpySparseDepthOutput(),
 849        invoker=_sparse_depth_invoker,
 850        name=name,
 851    )
 852
 853
 854# --------------------------------------------------------------------------- #
 855# --------------------------------------------------------------------------- #
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866# --------------------------------------------------------------------------- #
 867# Keypoint matching — fn(rgb_a, rgb_b) -> {points0, points1}
 868# --------------------------------------------------------------------------- #
 869
 870
 871class _NumpyKeypointInput:
 872    """Keypoint matching input: rgb_a + rgb_b from sample metadata."""
 873
 874    def prepare(self, sample: Sample) -> PreparedInput:
 875        rgb_a = np.asarray(sample.rgb, dtype=np.uint8)
 876        rgb_b = None
 877        if sample.metadata is not None:
 878            rgb_b = sample.metadata.get("rgb_b")
 879        if rgb_b is None:
 880            raise AdapterError(
 881                "Keypoint matching numpy adapter requires a second RGB "
 882                "frame (`rgb_b`) loaded by the dataset.",
 883                hint="Check that your manifest entries include `rgb_b` alongside `rgb`.",
 884            )
 885        return PreparedInput(
 886            payload={"rgb_a": rgb_a, "rgb_b": np.asarray(rgb_b, dtype=np.uint8)},
 887            context={"target_hw": rgb_a.shape[:2]},
 888        )
 889
 890
 891class _NumpyKeypointOutput:
 892    def finalize(
 893        self,
 894        model_output: Any,
 895        context: Dict[str, Any],
 896        sample: Sample,
 897    ) -> KeypointCorrespondencePrediction:
 898        if isinstance(model_output, dict):
 899            p0 = np.asarray(model_output.get("points0"), dtype=np.float32)
 900            p1 = np.asarray(model_output.get("points1"), dtype=np.float32)
 901            scores = model_output.get("scores")
 902        elif isinstance(model_output, tuple) and len(model_output) in (2, 3):
 903            p0 = np.asarray(model_output[0], dtype=np.float32)
 904            p1 = np.asarray(model_output[1], dtype=np.float32)
 905            scores = model_output[2] if len(model_output) == 3 else None
 906        else:
 907            raise AdapterError(
 908                "keypoint matching model must return a dict with keys "
 909                "{points0, points1[, scores]} or a 2/3-tuple.",
 910            )
 911        scores_arr = (
 912            np.asarray(scores, dtype=np.float32).reshape(-1) if scores is not None else None
 913        )
 914        return KeypointCorrespondencePrediction(
 915            points0=p0.reshape(-1, 2),
 916            points1=p1.reshape(-1, 2),
 917            scores=scores_arr,
 918        )
 919
 920
 921def _keypoint_invoker(model: Any, payload: Dict[str, Any]) -> Any:
 922    return model(payload["rgb_a"], payload["rgb_b"])
 923
 924
 925def make_numpy_keypoint_model(
 926    fn: Callable[[np.ndarray, np.ndarray], Any],
 927    *,
 928    name: str = "numpy_keypoint_model",
 929) -> BenchmarkableModel:
 930    """Wrap a keypoint-matching callable as a :class:`BenchmarkableModel`.
 931
 932    The callable takes ``(rgb_a, rgb_b)`` and returns either a dict
 933    with keys ``"points0"``, ``"points1"`` and optional ``"scores"``
 934    or a 2/3-tuple in the same order. Points are ``(N, 2)`` pixel
 935    coordinates.
 936
 937    Examples
 938    --------
 939    >>> import numpy as np
 940    >>> import rpx_benchmark as rpx
 941    >>> def my_matcher(rgb_a, rgb_b):
 942    ...     pts = np.array([[10, 10], [20, 20]], dtype=np.float32)
 943    ...     return pts, pts
 944    >>> bm = rpx.make_numpy_keypoint_model(my_matcher)
 945    """
 946    return BenchmarkableModel(
 947        task=TaskType.KEYPOINT_MATCHING,
 948        input_adapter=_NumpyKeypointInput(),
 949        model=fn,
 950        output_adapter=_NumpyKeypointOutput(),
 951        invoker=_keypoint_invoker,
 952        name=name,
 953    )
 954
 955
 956# --------------------------------------------------------------------------- #
 957# Object tracking — fn(rgb) -> [{track_id, boxes, scores?}, ...]
 958# --------------------------------------------------------------------------- #
 959
 960
 961class _NumpyTrackingInput:
 962    """Tracking input: a single RGB frame per step.
 963
 964    Per-frame tracking protocol: the model is handed one frame and
 965    must return the *currently active* tracks with a persistent
 966    ``track_id`` per track. The surrounding benchmark pipeline is
 967    responsible for accumulating tracks into a full :class:`Tracklet`
 968    set at scene boundaries.
 969    """
 970
 971    def prepare(self, sample: Sample) -> PreparedInput:
 972        rgb = np.asarray(sample.rgb, dtype=np.uint8)
 973        return PreparedInput(
 974            payload={"rgb": rgb},
 975            context={"target_hw": rgb.shape[:2], "sample_id": sample.id},
 976        )
 977
 978
 979class _NumpyTrackingOutput:
 980    def finalize(
 981        self,
 982        model_output: Any,
 983        context: Dict[str, Any],
 984        sample: Sample,
 985    ) -> TrackletPrediction:
 986        tracks = _coerce_tracking_output(model_output)
 987        return TrackletPrediction(tracks=tracks)
 988
 989
 990def _coerce_tracking_output(model_output: Any) -> list[Tracklet]:
 991    """Accept a variety of tracking return shapes; normalise to Tracklets.
 992
 993    Supported shapes:
 994        * ``TrackletPrediction(tracks=[...])`` — returned verbatim.
 995        * ``[Tracklet, ...]`` — returned verbatim.
 996        * ``[{"track_id": str, "boxes": (T, 4), "scores"?: (T,)}, ...]``
 997        * ``{"tracks": [as above]}`` — unwrapped.
 998    """
 999    if isinstance(model_output, TrackletPrediction):
1000        return list(model_output.tracks)
1001    if isinstance(model_output, dict) and "tracks" in model_output:
1002        model_output = model_output["tracks"]
1003    if not isinstance(model_output, Sequence):
1004        raise AdapterError(
1005            "Tracking model must return a sequence of tracks or a TrackletPrediction.",
1006            hint="Return either [Tracklet, ...] or "
1007            "[{'track_id': str, 'boxes': (T,4) array, ...}, ...].",
1008        )
1009    tracks: list[Tracklet] = []
1010    for item in model_output:
1011        if isinstance(item, Tracklet):
1012            tracks.append(item)
1013            continue
1014        if not isinstance(item, dict):
1015            raise AdapterError(
1016                f"Tracking output entries must be Tracklet or dict, got {type(item).__name__}.",
1017            )
1018        try:
1019            boxes = np.asarray(item["boxes"], dtype=np.float32).reshape(-1, 4)
1020        except (KeyError, ValueError) as e:
1021            raise AdapterError(
1022                f"Tracklet dict is missing / malformed 'boxes': {e}",
1023                hint="Boxes must be a (T, 4) array of [x1,y1,x2,y2] floats.",
1024            ) from e
1025        scores_raw = item.get("scores")
1026        scores = (
1027            np.asarray(scores_raw, dtype=np.float32).reshape(-1) if scores_raw is not None else None
1028        )
1029        tracks.append(
1030            Tracklet(
1031                track_id=str(item.get("track_id", f"t{len(tracks)}")),
1032                boxes=boxes,
1033                scores=scores,
1034            )
1035        )
1036    return tracks
1037
1038
1039def _tracking_invoker(model: Any, payload: Dict[str, Any]) -> Any:
1040    return model(payload["rgb"])
1041
1042
1043def make_numpy_tracking_model(
1044    fn: Callable[[np.ndarray], Any],
1045    *,
1046    name: str = "numpy_tracking_model",
1047) -> BenchmarkableModel:
1048    """Wrap a per-frame tracking callable as a :class:`BenchmarkableModel`.
1049
1050    The callable takes a single ``rgb`` ``H×W×3 uint8`` frame and
1051    returns the active tracks. Accepted return shapes (see
1052    :func:`_coerce_tracking_output`):
1053
1054    - ``TrackletPrediction``
1055    - list of :class:`Tracklet`
1056    - list of ``{"track_id": str, "boxes": (T, 4), "scores"?: (T,)}``
1057    - ``{"tracks": [as above]}``
1058
1059    Examples
1060    --------
1061    >>> import numpy as np
1062    >>> import rpx_benchmark as rpx
1063    >>> def my_tracker(rgb):
1064    ...     return [{
1065    ...         "track_id": "obj_0",
1066    ...         "boxes": np.array([[10, 10, 50, 50]], dtype=np.float32),
1067    ...     }]
1068    >>> bm = rpx.make_numpy_tracking_model(my_tracker)
1069    """
1070    return BenchmarkableModel(
1071        task=TaskType.OBJECT_TRACKING,
1072        input_adapter=_NumpyTrackingInput(),
1073        model=fn,
1074        output_adapter=_NumpyTrackingOutput(),
1075        invoker=_tracking_invoker,
1076        name=name,
1077    )