rpx_benchmark.data.hf_bridge

  1"""Row-level bridge between ``datasets.Dataset`` and :class:`Sample`.
  2
  3This module is the only place in the codebase that knows both the HF
  4schema (from :mod:`rpx_benchmark.data.features`) and the RPX in-memory
  5sample contract (from :mod:`rpx_benchmark.api`). Every other consumer
  6— the loader, the runner, adapters, metrics — sees only ``Sample``.
  7
  8Two surfaces:
  9
 10- :func:`row_to_sample` turns one HF row into a :class:`Sample` with
 11  the correct task-specific ground-truth dataclass attached.
 12- :class:`RPXHFBridge` wraps an iterable ``datasets.Dataset`` and
 13  yields :class:`Sample` batches shaped like
 14  :class:`rpx_benchmark.loader.RPXDataset`, so the runner can consume
 15  either interchangeably.
 16"""
 17
 18from __future__ import annotations
 19
 20from dataclasses import dataclass
 21from typing import Any, Dict, Iterator, List
 22from typing import Sequence as TSequence
 23
 24import numpy as np
 25
 26from ..api import (
 27    DepthGroundTruth,
 28    DetectionGroundTruth,
 29    Difficulty,
 30    KeypointCorrespondenceGroundTruth,
 31    Phase,
 32    RelativePoseGroundTruth,
 33    Sample,
 34    SegmentationGroundTruth,
 35    SparseDepthGroundTruth,
 36    TaskType,
 37    Tracklet,
 38    TrackletGroundTruth,
 39    VisualGroundingGroundTruth,
 40)
 41from ..exceptions import ManifestError
 42
 43__all__ = ["row_to_sample", "RPXHFBridge"]
 44
 45
 46# --------------------------------------------------------------------------- #
 47# Helpers
 48# --------------------------------------------------------------------------- #
 49
 50
 51def _decode_image(value: Any) -> np.ndarray:
 52    """Decode a HF ``Image`` column value into a numpy array.
 53
 54    ``datasets`` yields PIL images when ``decode=True``; when a user
 55    builds a Dataset manually the column may already be a numpy array.
 56    Both are supported.
 57    """
 58    if value is None:
 59        raise ManifestError("Image column is None; row is missing modality.")
 60    if isinstance(value, np.ndarray):
 61        return value
 62    # PIL.Image import deferred so the bridge doesn't force a PIL
 63    # dependency at module-load time.
 64    try:
 65        from PIL import Image as PILImage  # noqa: PLC0415
 66    except ImportError as e:  # pragma: no cover
 67        raise ImportError(
 68            "Decoding HF Image features requires Pillow. Install with: pip install Pillow"
 69        ) from e
 70    if isinstance(value, PILImage.Image):
 71        return np.asarray(value)
 72    # `datasets` can yield a dict {"bytes": ..., "path": ...} if the
 73    # row was fetched before decode. Redirect through PIL.
 74    if isinstance(value, dict) and "bytes" in value:
 75        import io  # noqa: PLC0415
 76
 77        return np.asarray(PILImage.open(io.BytesIO(value["bytes"])))
 78    raise ManifestError(f"Unsupported image-column value type: {type(value).__name__}")
 79
 80
 81def _decode_depth(value: Any) -> np.ndarray:
 82    """Depth arrives as a 16-bit PNG in millimetres; return float32 metres.
 83
 84    Fault-proof contract: the underlying array must be ``uint16``. A
 85    ``uint8`` result here would silently truncate every depth value to
 86    8-bit resolution — exactly the failure mode the dataset_hub's
 87    fault-proof guarantees exist to prevent.
 88    """
 89    arr = _decode_image(value)
 90    if arr.dtype != np.uint16:
 91        raise ManifestError(
 92            f"Depth decode contract violated: expected uint16, got dtype {arr.dtype}. "
 93            f"A uint8 result here silently truncates every depth value "
 94            f"to 8-bit resolution.",
 95            hint="RPX depth is stored as a single-channel 16-bit PNG in "
 96            "millimetres. Re-check the dataset's shard-generation script.",
 97        )
 98    if arr.ndim != 2:
 99        raise ManifestError(
100            f"Depth map must be 2-D, got shape {arr.shape}.",
101            hint="RPX depth is stored as a single-channel 16-bit PNG in "
102            "millimetres. Check the dataset's shard-generation script.",
103        )
104    depth = arr.astype(np.float32) / 1000.0
105    depth[arr == 0] = 0.0
106    return depth
107
108
109def _decode_mask(value: Any) -> np.ndarray:
110    """Instance mask as int32 H×W with pixel values = instance IDs.
111
112    Fault-proof contract: must end up 2-D ``int32`` so palette / I;16 /
113    L modes all normalise to the same integer-ID representation.
114    """
115    arr = _decode_image(value)
116    if arr.ndim != 2:
117        raise ManifestError(
118            f"Instance mask must be 2-D, got shape {arr.shape}.",
119        )
120    return arr.astype(np.int32)
121
122
123def _unflatten_pose(value: TSequence[float] | None) -> np.ndarray | None:
124    """Turn a flattened 4×4 pose (16 floats) into a numpy SE(3) matrix."""
125    if value is None:
126        return None
127    arr = np.asarray(list(value), dtype=np.float64)
128    if arr.size != 16:
129        raise ManifestError(
130            f"Pose must be a flat 16-element SE(3), got size {arr.size}.",
131        )
132    return arr.reshape(4, 4)
133
134
135def _enum_or_none(cls: Any, raw: Any) -> Any:
136    if raw is None or raw == "":
137        return None
138    return cls(raw)
139
140
141# --------------------------------------------------------------------------- #
142# Task-specific ground truth builders
143# --------------------------------------------------------------------------- #
144
145
146def _gt_depth(row: Dict[str, Any]) -> DepthGroundTruth:
147    return DepthGroundTruth(depth_map=_decode_depth(row["depth"]))
148
149
150def _gt_detection(row: Dict[str, Any]) -> DetectionGroundTruth:
151    boxes = np.asarray(row.get("boxes") or [], dtype=np.float32).reshape(-1, 4)
152    labels = list(row.get("labels") or [])
153    return DetectionGroundTruth(boxes=boxes, labels=labels)
154
155
156def _gt_segmentation(row: Dict[str, Any]) -> SegmentationGroundTruth:
157    return SegmentationGroundTruth(mask=_decode_mask(row["mask"]))
158
159
160def _gt_tracking(row: Dict[str, Any]) -> TrackletGroundTruth:
161    tracks_raw = row.get("tracks") or []
162    tracks: List[Tracklet] = []
163    for item in tracks_raw:
164        boxes = np.asarray(item["boxes"], dtype=np.float32).reshape(-1, 4)
165        scores_raw = item.get("scores")
166        scores = np.asarray(scores_raw, dtype=np.float32) if scores_raw else None
167        tracks.append(Tracklet(track_id=str(item["track_id"]), boxes=boxes, scores=scores))
168    return TrackletGroundTruth(tracks=tracks)
169
170
171def _gt_relative_pose(row: Dict[str, Any]) -> RelativePoseGroundTruth:
172    pose_a = _unflatten_pose(row["pose_a"])
173    pose_b = _unflatten_pose(row["pose_b"])
174    assert pose_a is not None and pose_b is not None  # schema guarantees both
175    rel = np.linalg.inv(pose_a) @ pose_b
176    return RelativePoseGroundTruth(rotation=rel[:3, :3], translation=rel[:3, 3])
177
178
179def _gt_visual_grounding(row: Dict[str, Any]) -> VisualGroundingGroundTruth:
180    boxes = np.asarray(row.get("boxes") or [], dtype=np.float32).reshape(-1, 4)
181    labels = list(row["labels"]) if row.get("labels") else None
182    return VisualGroundingGroundTruth(text=row["text"], boxes=boxes, labels=labels)
183
184
185def _gt_sparse_depth(row: Dict[str, Any]) -> SparseDepthGroundTruth:
186    coords = np.asarray(row["coordinates"], dtype=np.float32).reshape(-1, 2)
187    depths = np.asarray(row["depths"], dtype=np.float32)
188    return SparseDepthGroundTruth(coordinates=coords, depths=depths)
189
190
191
192
193def _gt_keypoints(row: Dict[str, Any]) -> KeypointCorrespondenceGroundTruth:
194    p0 = np.asarray(row["points0"], dtype=np.float32).reshape(-1, 2)
195    p1 = np.asarray(row["points1"], dtype=np.float32).reshape(-1, 2)
196    vis_raw = row.get("visibility")
197    vis = np.asarray(vis_raw, dtype=bool) if vis_raw is not None else None
198    return KeypointCorrespondenceGroundTruth(points0=p0, points1=p1, visibility=vis)
199
200
201_GT_BUILDERS = {
202    TaskType.MONOCULAR_DEPTH: _gt_depth,
203    TaskType.OBJECT_DETECTION: _gt_detection,
204    TaskType.OPEN_VOCAB_DETECTION: _gt_detection,
205    TaskType.OBJECT_SEGMENTATION: _gt_segmentation,
206    TaskType.OBJECT_TRACKING: _gt_tracking,
207    TaskType.RELATIVE_CAMERA_POSE: _gt_relative_pose,
208    TaskType.VISUAL_GROUNDING: _gt_visual_grounding,
209    TaskType.SPARSE_DEPTH: _gt_sparse_depth,
210    TaskType.KEYPOINT_MATCHING: _gt_keypoints,
211}
212
213
214# --------------------------------------------------------------------------- #
215# Public API
216# --------------------------------------------------------------------------- #
217
218
219def row_to_sample(row: Dict[str, Any], task: TaskType | str) -> Sample:
220    """Convert one HF ``Dataset`` row into a :class:`Sample`.
221
222    Parameters
223    ----------
224    row : dict
225        Single row fetched from ``hf_ds[i]``. Column layout must match
226        :func:`rpx_benchmark.data.features.features_for_task`.
227    task : TaskType or str
228
229    Returns
230    -------
231    Sample
232
233    Raises
234    ------
235    ManifestError
236        If required columns are missing or carry invalid shapes.
237    """
238    if isinstance(task, str):
239        task = TaskType(task)
240
241    metadata: Dict[str, Any] = {}
242    # Sibling RGB for pair tasks lives in ``rgb_b``; we preserve it so
243    # downstream adapters can see it without re-reading the dataset.
244    if row.get("rgb_b") is not None:
245        metadata["rgb_b"] = _decode_image(row["rgb_b"])
246
247    gt_builder = _GT_BUILDERS.get(task)
248    if gt_builder is None:
249        raise ManifestError(f"No HF bridge registered for task {task}.")
250    gt = gt_builder(row)
251
252    return Sample(
253        id=str(row["id"]),
254        rgb=_decode_image(row["rgb"]),
255        ground_truth=gt,
256        metadata=metadata or None,
257        phase=_enum_or_none(Phase, row.get("phase")),
258        difficulty=_enum_or_none(Difficulty, row.get("difficulty")),
259        camera_pose=_unflatten_pose(row.get("camera_pose")),
260    )
261
262
263@dataclass
264class RPXHFBridge:
265    """Iterable adapter: ``datasets.Dataset`` → :class:`Sample` batches.
266
267    Shaped like :class:`rpx_benchmark.loader.RPXDataset` (``__len__`` +
268    ``__iter__`` yielding ``list[Sample]``) so the runner can consume
269    either source interchangeably.
270
271    Parameters
272    ----------
273    hf_dataset : datasets.Dataset
274        A dataset with columns matching the canonical task schema.
275    task : TaskType
276    batch_size : int, default 1
277    """
278
279    hf_dataset: Any
280    task: TaskType
281    batch_size: int = 1
282
283    def __len__(self) -> int:
284        return int(len(self.hf_dataset))
285
286    def __iter__(self) -> Iterator[List[Sample]]:
287        batch: List[Sample] = []
288        for row in self.hf_dataset:
289            batch.append(row_to_sample(row, self.task))
290            if len(batch) == self.batch_size:
291                yield batch
292                batch = []
293        if batch:
294            yield batch