rpx_benchmark.loader

  1"""Data loading utilities for RPX benchmark.
  2
  3Reads a JSON *manifest* describing a slice of the RPX dataset (usually
  4produced by :func:`rpx_benchmark.hub.download_split`) and yields
  5:class:`~rpx_benchmark.api.Sample` batches the :class:`BenchmarkRunner`
  6can feed to a model.
  7
  8Manifest format
  9---------------
 10
 11::
 12
 13    {
 14      "task":   "monocular_depth",
 15      "root":   "/absolute/path/or/snapshot/root",
 16      "samples": [
 17        {
 18          "id":         "scene_001_clutter_00000",
 19          "scene":      "scene_001",
 20          "phase":      "clutter",
 21          "difficulty": "hard",
 22          "rgb":   "scenes/scene_001/0/rgb/00000.png",
 23          "depth": "scenes/scene_001/0/depth/00000.png",
 24          "mask":  "scenes/scene_001/0/mask/00000.png",
 25          "pose":  "scenes/scene_001/0/pose/00000.npz"
 26        },
 27        ...
 28      ]
 29    }
 30
 31All sample paths are resolved relative to ``root`` unless they are
 32absolute. ``depth`` files are 16-bit PNGs in millimetres (as saved by
 33``save_device_data.py``); the loader converts to float32 metres on
 34read. ``pose`` files are ``.npz`` with ``position`` and ``orientation``
 35(T265 convention, ``[x, y, z, w]`` quaternion).
 36
 37Errors
 38------
 39
 40All load-time failures are raised as
 41:class:`~rpx_benchmark.exceptions.ManifestError` so user code can
 42catch them specifically.
 43"""
 44
 45from __future__ import annotations
 46
 47import json
 48from dataclasses import dataclass
 49from pathlib import Path
 50from typing import Any, Dict, Iterator, List
 51
 52import numpy as np
 53
 54from .api import (
 55    DepthGroundTruth,
 56    DetectionGroundTruth,
 57    Difficulty,
 58    KeypointCorrespondenceGroundTruth,
 59    Phase,
 60    RelativePoseGroundTruth,
 61    Sample,
 62    SegmentationGroundTruth,
 63    SparseDepthGroundTruth,
 64    TaskType,
 65    Tracklet,
 66    TrackletGroundTruth,
 67    VisualGroundingGroundTruth,
 68)
 69from .exceptions import ManifestError
 70from .logging_utils import get_logger
 71
 72log = get_logger(__name__)
 73
 74# Phase index → Phase enum (matches capture session folder 0/1/2)
 75_PHASE_INDEX: Dict[int, Phase] = {
 76    0: Phase.CLUTTER,
 77    1: Phase.INTERACTION,
 78    2: Phase.CLEAN,
 79}
 80
 81
 82@dataclass
 83class RPXDataset:
 84    """Iterates over RPX samples for a specific task.
 85
 86    Manifest format (JSON)::
 87
 88        {
 89          "task": "object_segmentation",
 90          "root": "/path/to/data",
 91          "samples": [
 92            {
 93              "id": "scene_001_clutter_00000",
 94              "scene": "scene_001",
 95              "phase": "clutter",
 96              "difficulty": "hard",
 97              "rgb":   "scene_001/0/rgb/00000.png",
 98              "depth": "scene_001/0/depth/00000.png",
 99              "mask":  "scene_001/0/mask/00000.png",
100              "pose":  "scene_001/0/pose/00000.npz",
101              ...
102            }
103          ]
104        }
105
106    All paths are relative to ``root``.
107    ``depth`` files are 16-bit PNG in millimetres (as saved by save_device_data.py).
108    ``pose`` files are NPZ with keys ``position`` ([x,y,z] metres) and
109    ``orientation`` ([x,y,z,w] quaternion) from the T265 tracker.
110    """
111
112    samples: List[Dict[str, Any]]
113    task: TaskType
114    root: Path
115    batch_size: int = 1
116
117    @classmethod
118    def from_manifest(
119        cls,
120        manifest_path: str | Path,
121        batch_size: int = 1,
122        validate: bool = False,
123        max_samples: int | None = None,
124    ) -> "RPXDataset":
125        """Load a manifest JSON file from disk and return a dataset.
126
127        Parameters
128        ----------
129        manifest_path : str or Path
130            Path to the manifest JSON. Produced either by
131            :func:`rpx_benchmark.hub.download_split` or by a custom
132            upload script.
133        batch_size : int
134            Number of samples per iteration. Default 1.
135        validate : bool, optional
136            When ``True``, validate the manifest against the strict
137            Pydantic schema in :mod:`rpx_benchmark.schemas` before
138            constructing the dataset. Requires the ``schemas`` extra
139            (``pip install 'rpx-benchmark[schemas]'``). Default
140            ``False`` preserves historical, tolerant behaviour.
141        max_samples : int, optional
142            Keep only the first N manifest samples. Intended for smoke
143            validation; must be at least 1 when provided.
144
145        Returns
146        -------
147        RPXDataset
148
149        Raises
150        ------
151        ManifestError
152            If the manifest file is missing, not valid JSON, is missing
153            required top-level fields, or (when ``validate=True``)
154            fails strict schema validation.
155        """
156        manifest_path = Path(manifest_path)
157        if not manifest_path.is_file():
158            raise ManifestError(
159                f"Manifest file not found: {manifest_path}",
160                hint="Did the HuggingFace download fail? Try rerunning with "
161                "--cache-dir pointing at a writable location.",
162            )
163        try:
164            with manifest_path.open("r", encoding="utf-8") as f:
165                manifest = json.load(f)
166        except json.JSONDecodeError as e:
167            raise ManifestError(
168                f"Manifest at {manifest_path} is not valid JSON: {e}",
169            ) from e
170        dataset = cls.from_dict(
171            manifest,
172            batch_size=batch_size,
173            default_root=manifest_path.parent,
174            validate=validate,
175        )
176        if max_samples is not None:
177            if max_samples < 1:
178                raise ManifestError(
179                    f"max_samples must be >= 1, got {max_samples}",
180                )
181            dataset.samples = dataset.samples[:max_samples]
182        return dataset
183
184    @classmethod
185    def from_dict(
186        cls,
187        manifest: Dict[str, Any],
188        batch_size: int = 1,
189        default_root: str | Path | None = None,
190        validate: bool = False,
191    ) -> "RPXDataset":
192        """Build a dataset from an already-parsed manifest dict.
193
194        Raises
195        ------
196        ManifestError
197            If ``task`` is missing or unknown, or if ``samples`` is
198            missing. When ``validate=True``, any Pydantic validation
199            error is also surfaced as :class:`ManifestError`.
200        """
201        if validate:
202            # Opt-in strict validation via rpx_benchmark.schemas. Kept
203            # lazy so users who don't install the ``schemas`` extra
204            # don't hit pydantic import at module load time.
205            from .schemas import validate_manifest  # noqa: PLC0415 — lazy import
206
207            model = validate_manifest(manifest)
208            manifest = model.model_dump(exclude_none=True)
209
210        if "task" not in manifest:
211            raise ManifestError(
212                "Manifest is missing required field 'task'.",
213                hint="Task must be one of: " + ", ".join(t.value for t in TaskType),
214            )
215        try:
216            task = TaskType(manifest["task"])
217        except ValueError as e:
218            raise ManifestError(
219                f"Manifest task {manifest['task']!r} is not a known TaskType.",
220                hint="Expected one of: " + ", ".join(t.value for t in TaskType),
221            ) from e
222
223        if "samples" not in manifest:
224            raise ManifestError(
225                "Manifest is missing required field 'samples'.",
226            )
227
228        root = Path(manifest.get("root") or default_root or ".")
229        samples = manifest["samples"]
230        if not isinstance(samples, list):
231            raise ManifestError(
232                f"Manifest 'samples' must be a list, got {type(samples).__name__}",
233            )
234        log.debug("loaded manifest: task=%s root=%s samples=%d", task.value, root, len(samples))
235        return cls(samples=samples, task=task, root=root, batch_size=batch_size)
236
237    @classmethod
238    def from_hf(
239        cls,
240        hf_dataset: Any,
241        task: TaskType | str | None = None,
242        batch_size: int = 1,
243    ) -> Any:
244        """Wrap a ``datasets.Dataset`` as an iterable RPX source.
245
246        Returns a :class:`rpx_benchmark.data.hf_bridge.RPXHFBridge`
247        which is interchangeable with :class:`RPXDataset` at the
248        runner's consumer interface (``__len__`` + ``__iter__`` over
249        ``list[Sample]``).
250
251        Parameters
252        ----------
253        hf_dataset : datasets.Dataset
254            Dataset whose columns match
255            :func:`rpx_benchmark.data.features.features_for_task`.
256        task : TaskType or str, optional
257            Task this dataset serves. If omitted, we try to infer it
258            from ``hf_dataset.info.description`` (where
259            :func:`rpx_benchmark.data.load_hf` stamps it).
260        batch_size : int, default 1
261
262        Returns
263        -------
264        RPXHFBridge
265
266        Raises
267        ------
268        ManifestError
269            If ``task`` is neither provided nor discoverable on the
270            dataset metadata.
271        """
272        from .data.hf_bridge import RPXHFBridge  # noqa: PLC0415 — lazy
273
274        if task is None:
275            info = getattr(hf_dataset, "info", None)
276            config = getattr(info, "config_name", None) if info is not None else None
277            if config is None:
278                raise ManifestError(
279                    "Could not infer task from the HF dataset; pass task=TaskType.* explicitly.",
280                )
281            task = TaskType(config)
282        if isinstance(task, str):
283            task = TaskType(task)
284        return RPXHFBridge(hf_dataset=hf_dataset, task=task, batch_size=batch_size)
285
286    def __len__(self) -> int:
287        return len(self.samples)
288
289    def __iter__(self) -> Iterator[List[Sample]]:
290        batch: List[Sample] = []
291        for entry in self.samples:
292            batch.append(self._load_sample(entry))
293            if len(batch) == self.batch_size:
294                yield batch
295                batch = []
296        if batch:
297            yield batch
298
299    def _load_sample(self, entry: Dict[str, Any]) -> Sample:
300        rgb = self._load_rgb(entry["rgb"])
301
302        # Phase — from string label or integer session index
303        phase: Phase | None = None
304        if "phase" in entry:
305            v = entry["phase"]
306            phase = Phase(v) if isinstance(v, str) else _PHASE_INDEX.get(int(v))
307
308        # Difficulty
309        difficulty: Difficulty | None = None
310        if "difficulty" in entry:
311            difficulty = Difficulty(entry["difficulty"])
312
313        # Camera pose (T265 NPZ → 4×4 SE(3) float64)
314        camera_pose: np.ndarray | None = None
315        if "pose" in entry:
316            camera_pose = self._load_pose(entry["pose"])
317
318        # Fisheye images — stored in metadata if present
319        metadata: Dict[str, Any] = dict(entry.get("metadata") or {})
320        scene = entry.get("scene_id") or entry.get("scene")
321        if scene is not None:
322            metadata["scene_id"] = scene
323        if "fisheye_left" in entry:
324            metadata["fisheye_left"] = self._load_gray(entry["fisheye_left"])
325        if "fisheye_right" in entry:
326            metadata["fisheye_right"] = self._load_gray(entry["fisheye_right"])
327        # Second RGB frame for pair tasks (relative pose, keypoint matching)
328        if "rgb_b" in entry:
329            metadata["rgb_b"] = self._load_rgb(entry["rgb_b"])
330
331        gt = self._load_ground_truth(entry)
332
333        return Sample(
334            id=entry["id"],
335            rgb=rgb,
336            ground_truth=gt,
337            metadata=metadata or None,
338            phase=phase,
339            difficulty=difficulty,
340            camera_pose=camera_pose,
341        )
342
343    def _load_ground_truth(self, entry: Dict[str, Any]) -> Any:
344        task = self.task
345
346        if task == TaskType.MONOCULAR_DEPTH:
347            return DepthGroundTruth(depth_map=self._load_depth(entry["depth"]))
348
349        if task in (TaskType.OBJECT_DETECTION, TaskType.OPEN_VOCAB_DETECTION):
350            boxes, labels = self._load_boxes(entry["boxes"])
351            return DetectionGroundTruth(boxes=boxes, labels=labels)
352
353        if task == TaskType.OBJECT_SEGMENTATION:
354            return SegmentationGroundTruth(mask=self._load_mask(entry["mask"]))
355
356        if task == TaskType.OBJECT_TRACKING:
357            if "tracks" in entry and self._resolve(entry["tracks"]).is_file():
358                return self._load_tracklets(entry)
359            return self._load_tracking_mask(entry)
360
361        if task == TaskType.RELATIVE_CAMERA_POSE:
362            return self._load_relative_pose(entry)
363
364        if task == TaskType.VISUAL_GROUNDING:
365            return self._load_visual_grounding(entry)
366
367        if task == TaskType.SPARSE_DEPTH:
368            return self._load_sparse_depth(entry)
369
370
371        if task == TaskType.KEYPOINT_MATCHING:
372            return self._load_keypoints(entry)
373
374        raise ManifestError(
375            f"Unsupported task {task} — loader does not know how to parse "
376            "its ground-truth entries.",
377        )
378
379    # ------------------------------------------------------------------ #
380    # Ground-truth loaders
381    # ------------------------------------------------------------------ #
382
383    def _load_tracklets(self, entry: Dict[str, Any]) -> TrackletGroundTruth:
384        path = self._resolve(entry["tracks"])
385        with path.open("r", encoding="utf-8") as f:
386            data = json.load(f)
387        tracks = []
388        for item in data:
389            boxes = np.array(item["boxes"], dtype=np.float32)
390            scores = np.array(item["scores"], dtype=np.float32) if "scores" in item else None
391            tracks.append(Tracklet(track_id=str(item["track_id"]), boxes=boxes, scores=scores))
392        return TrackletGroundTruth(tracks=tracks)
393
394    def _load_tracking_mask(self, entry: Dict[str, Any]) -> TrackletGroundTruth:
395        """Represent one released instance-mask frame as one-step tracklets.
396
397        The D3 production runner groups the masks into complete clips before
398        evaluation.  This fallback keeps the generic frame loader usable for
399        inspection without depending on the unpublished ``tracklets/v1.json``.
400        """
401        mask = self._load_mask(entry["mask"])
402        tracks = []
403        for track_id in np.unique(mask):
404            if int(track_id) <= 0:
405                continue
406            ys, xs = np.nonzero(mask == track_id)
407            if xs.size == 0:
408                continue
409            boxes = np.array(
410                [[xs.min(), ys.min(), xs.max() + 1, ys.max() + 1]],
411                dtype=np.float32,
412            )
413            tracks.append(Tracklet(track_id=str(int(track_id)), boxes=boxes))
414        return TrackletGroundTruth(tracks=tracks)
415
416    def _load_relative_pose(self, entry: Dict[str, Any]) -> RelativePoseGroundTruth:
417        from .pose_conventions import relative_pose_from_raw_t265
418
419        # Published RPX poses retain librealsense's T265 axes. RCPE models
420        # consume D435 RGB and conventionally emit OpenCV camera coordinates,
421        # so convert the fixed axis basis before forming the relative pose.
422        pose_a = self._load_pose(entry["pose_a"])  # 4×4 SE(3)
423        pose_b = self._load_pose(entry["pose_b"])
424        T_rel = relative_pose_from_raw_t265(pose_a, pose_b)
425        return RelativePoseGroundTruth(
426            rotation=T_rel[:3, :3],
427            translation=T_rel[:3, 3],
428        )
429
430    def _load_visual_grounding(self, entry: Dict[str, Any]) -> VisualGroundingGroundTruth:
431        text = entry["text"]
432        boxes_raw = entry.get("boxes")
433        labels = entry.get("labels")
434        if boxes_raw is None:
435            boxes = np.zeros((0, 4), dtype=np.float32)
436        elif isinstance(boxes_raw, str):
437            gt_boxes, _ = self._load_boxes(boxes_raw)
438            boxes = gt_boxes
439        else:
440            boxes = np.array(boxes_raw, dtype=np.float32)
441        return VisualGroundingGroundTruth(text=text, boxes=boxes, labels=labels)
442
443    def _load_sparse_depth(self, entry: Dict[str, Any]) -> SparseDepthGroundTruth:
444        coordinates = self._load_array_or_inline(entry["coordinates"]).astype(np.float32)
445        depths = self._load_array_or_inline(entry["depths"]).astype(np.float32)
446        return SparseDepthGroundTruth(coordinates=coordinates, depths=depths)
447
448
449    def _load_keypoints(self, entry: Dict[str, Any]) -> KeypointCorrespondenceGroundTruth:
450        points0 = self._load_array_or_inline(entry["points0"]).astype(np.float32)
451        points1 = self._load_array_or_inline(entry["points1"]).astype(np.float32)
452        visibility = None
453        if "visibility" in entry:
454            visibility = self._load_array_or_inline(entry["visibility"]).astype(bool)
455        return KeypointCorrespondenceGroundTruth(
456            points0=points0, points1=points1, visibility=visibility
457        )
458
459    # ------------------------------------------------------------------ #
460    # Low-level file loaders
461    # ------------------------------------------------------------------ #
462
463    def _resolve(self, relative_path: str | Path) -> Path:
464        path = Path(relative_path)
465        return path if path.is_absolute() else self.root / path
466
467    def _load_rgb(self, relative_path: str) -> np.ndarray:
468        """Load an RGB image → ``(H, W, 3) uint8``.
469
470        Delegates to :func:`rpx_benchmark.decode_contracts.safe_load_rgb`
471        so the dtype/shape contract is enforced in exactly one place,
472        shared with every adapter that decodes directly.
473        """
474        from .decode_contracts import safe_load_rgb
475
476        return safe_load_rgb(self._resolve(relative_path))
477
478    def _load_depth(self, relative_path: str) -> np.ndarray:
479        """Load a 16-bit PNG depth file → ``(H, W) float32`` in metres.
480
481        Decoding is delegated to
482        :func:`rpx_benchmark.decode_contracts.safe_load_depth`, which
483        guarantees ``uint16`` (the only correct decode for 16-bit
484        depth PNGs). We then cast to float32 and scale mm → m.
485        """
486        from .decode_contracts import safe_load_depth
487
488        raw = safe_load_depth(self._resolve(relative_path))  # (H, W) uint16
489        depth_mm = raw.astype(np.float32)
490        depth_m = depth_mm / 1000.0
491        depth_m[depth_mm == 0] = 0.0  # keep 0 as "invalid" sentinel
492        return depth_m
493
494    def _load_pose(self, relative_path: str) -> np.ndarray:
495        """Load a T265 pose → 4×4 SE(3) float64 (camera-to-world).
496
497        Two on-disk formats are accepted:
498
499        * Legacy ``.npz`` (per-frame zip): ``position`` (3,) +
500          ``orientation`` (4, xyzw quaternion).
501        * Compact ``.npy`` (per-frame raw, produced by the
502          ``lossless-convert`` step): a single ``(7,) float64`` vector
503          packing ``[x, y, z, qx, qy, qz, qw]``. The values are
504          bit-identical to the legacy ``.npz`` form.
505        """
506        path = self._resolve(relative_path)
507        data = np.load(path)
508        if path.suffix.lower() == ".npy":
509            arr = np.asarray(data, dtype=np.float64)
510            if arr.shape != (7,):
511                raise ManifestError(
512                    f"Pose .npy at {path} must be shape (7,) float64; got {arr.shape}",
513                    hint="The dataset_hub lossless-convert step writes a "
514                    "single (7,) vector packing [x, y, z, qx, qy, qz, qw].",
515                )
516            position = arr[:3]
517            quat_xyzw = arr[3:]
518        else:
519            position = data["position"].astype(np.float64)  # (3,)
520            quat_xyzw = data["orientation"].astype(np.float64)  # (4,) xyzw
521
522        R = _quat_xyzw_to_rotmat(quat_xyzw)
523        T = np.eye(4, dtype=np.float64)
524        T[:3, :3] = R
525        T[:3, 3] = position
526        if T.shape != (4, 4) or T.dtype != np.float64:
527            raise ManifestError(
528                f"Pose load contract violated at {path}: "
529                f"expected (4, 4) float64, got shape {T.shape} dtype {T.dtype}",
530            )
531        return T
532
533    def _load_mask(self, relative_path: str) -> np.ndarray:
534        """Load an instance-segmentation mask → ``(H, W) int32``.
535
536        Delegates to :func:`rpx_benchmark.decode_contracts.safe_load_mask`.
537        """
538        from .decode_contracts import safe_load_mask
539
540        return safe_load_mask(self._resolve(relative_path))
541
542    def _load_gray(self, relative_path: str) -> np.ndarray:
543        """Load a grayscale image (fisheye) → ``(H, W) uint8``.
544
545        Delegates to :func:`rpx_benchmark.decode_contracts.safe_load_gray`.
546        """
547        from .decode_contracts import safe_load_gray
548
549        return safe_load_gray(self._resolve(relative_path))
550
551    def _load_boxes(self, relative_path: str):
552        path = self._resolve(relative_path)
553        with path.open("r", encoding="utf-8") as f:
554            data = json.load(f)
555        boxes = np.array([item["bbox"] for item in data], dtype=np.float32)
556        labels = [item["label"] for item in data]
557        return boxes, labels
558
559    def _load_array_or_inline(self, value: Any) -> np.ndarray:
560        if isinstance(value, str):
561            path = self._resolve(value)
562            if path.suffix.lower() == ".npy":
563                return np.load(path)
564            with path.open("r", encoding="utf-8") as f:
565                return np.array(json.load(f))
566        return np.array(value)
567
568
569# ------------------------------------------------------------------ #
570# Quaternion helper
571# ------------------------------------------------------------------ #
572
573
574def _quat_xyzw_to_rotmat(q: np.ndarray) -> np.ndarray:
575    """Convert quaternion [x, y, z, w] (T265 convention) to 3×3 rotation matrix."""
576    x, y, z, w = q / np.linalg.norm(q)
577    return np.array(
578        [
579            [1 - 2 * y * y - 2 * z * z, 2 * x * y - 2 * z * w, 2 * x * z + 2 * y * w],
580            [2 * x * y + 2 * z * w, 1 - 2 * x * x - 2 * z * z, 2 * y * z - 2 * x * w],
581            [2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x * x - 2 * y * y],
582        ],
583        dtype=np.float64,
584    )