rpx_benchmark.video_loader

  1"""Sequence-shaped loader for video tasks (Video Depth today; future video tasks).
  2
  3Iterates per ``(scene, phase)`` clip, yielding
  4:class:`~rpx_benchmark.api.VideoSample`. The loader supports an
  5optional ``frame_budget`` + ``sampling`` mode for the temporal-
  6resolution ablation study (run at full 250 frames + downsampled
  7budgets such as 150/75/25 to characterise how each video depth model
  8degrades with sparser input).
  9
 10Cell-log granularity
 11--------------------
 12
 13One ``(scene, phase)`` clip → one cell-log row keyed
 14``(model, task, scene, phase, frame_budget)``. ``frame_budget=None``
 15(or the literal 0) is the canonical "full clip" row that feeds Φ;
 16non-None budgets feed the ablation plot. Φ is computed only over
 17``frame_budget=None`` rows to keep the comparison across models
 18honest.
 19
 20Sampling modes
 21--------------
 22
 23* ``"all"`` — every frame on disk (the default). ``frame_budget`` must
 24  be ``None`` in this mode.
 25* ``"stride"`` — uniform stride keeping the first and last frame
 26  anchored. Reduces to FPS-in-1D-timestamp when frames are evenly
 27  spaced (our 30 fps capture is).
 28* ``"fps_se3"`` — farthest-point sampling in 3D camera-translation
 29  space (T265 positions). First and last frame anchored. Selects
 30  ``frame_budget`` frames that maximise spatial coverage of the
 31  camera path. This is the most interesting mode for a robotics
 32  paper: it answers "does losing frames hurt because of less data
 33  or because we miss key viewpoints?" — only the second is
 34  diagnostic. Deterministic given ``(scene, phase, frame_budget)``.
 35
 36Per-clip scale-and-shift alignment for affine-invariant models is
 37applied at the runner level, not here — this loader is decode-only.
 38
 39Memory budget per full-budget clip at D435 resolution: ~540 MB
 40(250 × 640 × 480 × 3 B RGB + 250 × 640 × 480 × 4 B depth). Workable
 41on a 64 GB host with one clip in flight per worker.
 42"""
 43
 44from __future__ import annotations
 45
 46import json
 47from dataclasses import dataclass
 48from pathlib import Path
 49from typing import Iterator, List, Literal, Optional
 50
 51import numpy as np
 52
 53from .api import TaskType, VideoDepthGroundTruth, VideoSample
 54from .decode_contracts import safe_load_depth, safe_load_pose, safe_load_rgb
 55from .exceptions import ConfigError, ManifestError
 56from .logging_utils import get_logger
 57from .metrics.depth_alignment import DEPTH_MAX_M, DEPTH_MIN_M
 58
 59log = get_logger(__name__)
 60
 61
 62SamplingMode = Literal["all", "stride", "fps_se3"]
 63
 64
 65@dataclass
 66class VideoDepthDataset:
 67    """Iterates per ``(scene, phase)`` clip for the video-depth task.
 68
 69    Yielded value is :class:`~rpx_benchmark.api.VideoSample` whose
 70    ``rgb_seq`` is a ``(T, H, W, 3) uint8`` stack and whose
 71    ``ground_truth`` is :class:`~rpx_benchmark.api.VideoDepthGroundTruth`
 72    with the same ``T``. When ``frame_budget`` is ``None`` (the
 73    canonical run), ``T`` is whatever the phase has on disk; otherwise
 74    ``T == frame_budget`` and ``metadata['frame_indices']`` records
 75    which original-phase indices were kept.
 76
 77    Parameters
 78    ----------
 79    samples : list[dict]
 80        Per-clip manifest entries (one dict per (scene, phase)).
 81        Produced by
 82        :func:`rpx_benchmark.dataset_hub.split_manifests.write_split_manifests`
 83        in its ``"video_depth"`` mode. Each dict carries ``scene_id``,
 84        ``phase``, ``frame_filenames`` (sorted list), and optionally
 85        ``pose_filenames``, ``intrinsics`` (3×3 list), ``difficulty``.
 86    task : TaskType
 87        Must be :data:`~rpx_benchmark.api.TaskType.VIDEO_DEPTH`.
 88    root : Path
 89        Snapshot root that ``frame_filenames`` and ``pose_filenames``
 90        are resolved against.
 91    batch_size : int
 92        Always 1 for Video Depth — clips are not stackable. Kept for API
 93        uniformity with :class:`~rpx_benchmark.loader.RPXDataset`.
 94    frame_budget : int or None
 95        If set, subsample each clip to exactly this many frames via
 96        ``sampling``. ``None`` means use every frame on disk.
 97    sampling : Literal["all", "stride", "fps_se3"]
 98        Subsampling strategy. ``"all"`` requires
 99        ``frame_budget is None``. ``"stride"`` is uniform-with-anchored-
100        endpoints. ``"fps_se3"`` is farthest-point sampling in T265
101        position space (requires ``pose_filenames`` in the manifest
102        entry).
103    """
104
105    samples: List[dict]
106    task: TaskType
107    root: Path
108    batch_size: int = 1
109    frame_budget: Optional[int] = None
110    sampling: SamplingMode = "all"
111    compute_fscore: bool = False
112
113    def __post_init__(self) -> None:
114        if self.sampling == "all" and self.frame_budget is not None:
115            raise ConfigError(
116                "sampling='all' is incompatible with a frame_budget",
117                hint="Drop the budget (None means take every frame) or "
118                "pick sampling='stride' or 'fps_se3'.",
119            )
120        if self.sampling != "all" and self.frame_budget is None:
121            raise ConfigError(
122                f"sampling={self.sampling!r} requires a frame_budget",
123                hint="Pass frame_budget=<int> or set sampling='all'.",
124            )
125        if self.frame_budget is not None and self.frame_budget < 2:
126            raise ConfigError(
127                f"frame_budget must be >= 2 (got {self.frame_budget})",
128                hint="Temporal metrics need at least two frames.",
129            )
130
131    # ------------------------------------------------------------------ #
132    # Iteration
133    # ------------------------------------------------------------------ #
134
135    def __len__(self) -> int:
136        return len(self.samples)
137
138    def __iter__(self) -> Iterator[List[VideoSample]]:
139        for group in self.samples:
140            yield [self._load_one_clip(group)]
141
142    # ------------------------------------------------------------------ #
143    # Per-clip loading
144    # ------------------------------------------------------------------ #
145
146    def _load_one_clip(self, group: dict) -> VideoSample:
147        scene = group["scene_id"]
148        phase = group["phase"]
149        frame_files = list(group["frame_filenames"])
150        pose_files = list(group.get("pose_filenames") or [])
151
152        if not frame_files:
153            raise ManifestError(
154                f"clip {scene}/{phase}: frame_filenames is empty",
155            )
156        if pose_files and len(pose_files) != len(frame_files):
157            raise ManifestError(
158                f"clip {scene}/{phase}: pose_filenames length "
159                f"{len(pose_files)} != frame_filenames length "
160                f"{len(frame_files)}",
161            )
162        if pose_files:
163            missing_poses = [p for p in pose_files if not (self.root / p).is_file()]
164            if missing_poses:
165                if self.sampling == "fps_se3":
166                    raise ManifestError(
167                        f"clip {scene}/{phase}: sampling='fps_se3' requires pose files, "
168                        f"but {len(missing_poses)} are absent (first: {missing_poses[0]})",
169                    )
170                log.warning(
171                    "clip %s/%s: manifest references %d unavailable pose files; "
172                    "continuing RGB-D-only with pose diagnostics disabled",
173                    scene,
174                    phase,
175                    len(missing_poses),
176                )
177                pose_files = []
178
179        # 1. Decide which frame indices to keep.
180        T_full = len(frame_files)
181        if self.frame_budget is None:
182            keep = np.arange(T_full, dtype=np.int32)
183        else:
184            if self.frame_budget > T_full:
185                raise ManifestError(
186                    f"clip {scene}/{phase}: frame_budget="
187                    f"{self.frame_budget} exceeds available frame count "
188                    f"{T_full}",
189                )
190            if self.sampling == "stride":
191                keep = _stride_sample(T_full, self.frame_budget)
192            elif self.sampling == "fps_se3":
193                if not pose_files:
194                    raise ManifestError(
195                        f"clip {scene}/{phase}: sampling='fps_se3' "
196                        f"requires pose_filenames in the manifest entry",
197                    )
198                positions = _read_translation_track(self.root, pose_files)
199                keep = _fps_se3(positions, self.frame_budget)
200            else:  # pragma: no cover — guarded by __post_init__
201                raise ConfigError(
202                    f"unknown sampling mode {self.sampling!r}",
203                    hint="Pick one of 'all', 'stride', 'fps_se3'.",
204                )
205
206        # 2. Decode RGB + depth + poses ONLY for the kept indices.
207        #    Reading 250 PNGs then throwing 200 away would waste 80% of
208        #    the I/O budget; kept-only decode is the only sane strategy
209        #    when the budget is small.
210        rgb_seq = np.stack([safe_load_rgb(self.root / frame_files[i]) for i in keep])
211        # Frame files come from the rgb manifest; depth files mirror
212        # the layout (rgb/<stem>.png ↔ depth/<stem>.png). Group dicts
213        # carry an explicit ``depth_filenames`` list to avoid relying on
214        # that convention; fall back if it's absent.
215        depth_files = group.get("depth_filenames") or [
216            _swap_modality(p, "rgb", "depth") for p in frame_files
217        ]
218        depth_u16 = np.stack([safe_load_depth(self.root / depth_files[i]) for i in keep])
219        depth_seq = depth_u16.astype(np.float32) / 1000.0
220        valid_seq = (
221            np.isfinite(depth_seq)
222            & (depth_seq > DEPTH_MIN_M)
223            & (depth_seq < DEPTH_MAX_M)
224        )
225
226        camera_pose_seq: Optional[np.ndarray] = None
227        if pose_files:
228            poses = []
229            for i in keep:
230                pos, quat = safe_load_pose(self.root / pose_files[i])
231                T_mat = _se3_from_pos_quat(pos, quat)
232                poses.append(T_mat)
233            camera_pose_seq = np.stack(poses)
234
235        gt = VideoDepthGroundTruth(
236            depth_map_seq=depth_seq,
237            valid_mask_seq=valid_seq,
238            frame_indices=keep.astype(np.int32),
239            compute_fscore=self.compute_fscore,
240        )
241        # Stash optional extras the temporal calculator may want on the
242        # GT object (the calculator reads them via getattr so absence is
243        # graceful).
244        if camera_pose_seq is not None:
245            object.__setattr__(gt, "poses", camera_pose_seq)
246        object.__setattr__(gt, "rgb_seq", rgb_seq)
247
248        metadata: dict = {
249            "scene_id": scene,
250            "phase_idx": phase,
251            "frame_indices": keep.tolist(),
252            "frame_budget": int(self.frame_budget) if self.frame_budget else 0,
253            "sampling": self.sampling,
254        }
255        if "intrinsics" in group:
256            metadata["intrinsics"] = np.asarray(group["intrinsics"], dtype=np.float32)
257
258        from .api import Difficulty, Phase
259
260        _PHASE_BY_INT = {0: Phase.CLUTTER, 1: Phase.INTERACTION, 2: Phase.CLEAN}
261        try:
262            phase_enum: Phase | None = (
263                Phase(phase) if isinstance(phase, str) else _PHASE_BY_INT.get(int(phase))
264            )
265        except (ValueError, TypeError):
266            phase_enum = None
267        difficulty_str = group.get("difficulty")
268        try:
269            difficulty_enum: Difficulty | None = (
270                Difficulty(difficulty_str) if difficulty_str else None
271            )
272        except ValueError:
273            difficulty_enum = None
274
275        return VideoSample(
276            id=f"{scene}_{phase}_budget{metadata['frame_budget']}",
277            rgb_seq=rgb_seq,
278            ground_truth=gt,
279            metadata=metadata,
280            phase=phase_enum,
281            difficulty=difficulty_enum,
282            camera_pose_seq=camera_pose_seq,
283        )
284
285    # ------------------------------------------------------------------ #
286    # Constructors
287    # ------------------------------------------------------------------ #
288
289    @classmethod
290    def from_manifest(
291        cls,
292        manifest_path: str | Path,
293        batch_size: int = 1,
294        frame_budget: Optional[int] = None,
295        sampling: SamplingMode = "all",
296        max_samples: Optional[int] = None,
297        compute_fscore: bool = False,
298    ) -> "VideoDepthDataset":
299        """Build a VideoDepthDataset from a per-clip JSON manifest.
300
301        Expected JSON shape::
302
303            {
304              "task": "video_depth",
305              "root": "/abs/path/or/snapshot/root",
306              "samples": [
307                {
308                  "scene_id": "scene_001",
309                  "phase": 0,
310                  "frame_filenames": ["scenes/scene_001/0/rgb/00000.png", ...],
311                  "depth_filenames": ["scenes/scene_001/0/depth/00000.png", ...],
312                  "pose_filenames":  ["scenes/scene_001/0/cam_pose/00000.npz", ...],
313                  "intrinsics": [[fx,0,cx],[0,fy,cy],[0,0,1]],
314                  "difficulty": "easy"
315                },
316                ...
317              ]
318            }
319
320        The writer for this shape lives in
321        :func:`rpx_benchmark.dataset_hub.split_manifests.write_split_manifests`
322        under the ``video_depth`` recipe.
323        """
324        manifest_path = Path(manifest_path)
325        if not manifest_path.is_file():
326            raise ManifestError(
327                f"Manifest file not found: {manifest_path}",
328            )
329        try:
330            payload = json.loads(manifest_path.read_text(encoding="utf-8"))
331        except json.JSONDecodeError as e:
332            raise ManifestError(
333                f"Manifest at {manifest_path} is not valid JSON: {e}",
334            ) from e
335        task_str = payload.get("task")
336        if task_str != "video_depth":
337            raise ManifestError(
338                f"VideoDepthDataset.from_manifest expected task='video_depth'; got {task_str!r}",
339            )
340        # Resolve the modality-file root. The manifest's ``root`` field
341        # is the canonical source — every frame_filename / depth_filename
342        # / pose_filename is resolved relative to it. If the manifest
343        # omits ``root`` (or sets it to None / ""), fall back to the
344        # manifest file's parent dir; warn loudly so the caller notices
345        # before a misleading FileNotFoundError downstream.
346        manifest_root = payload.get("root")
347        if manifest_root:
348            root = Path(manifest_root)
349        else:
350            from .logging_utils import get_logger
351
352            get_logger(__name__).warning(
353                "VideoDepthDataset.from_manifest: %s has no 'root' field; falling "
354                "back to manifest parent dir %s. If your frame_filenames "
355                "are relative to a different root (e.g. an HF cache "
356                "snapshot or a staged dataset tree), set the 'root' field "
357                "explicitly in the manifest. See "
358                "benchmark/README.md for the manifest schema.",
359                manifest_path,
360                manifest_path.parent,
361            )
362            root = manifest_path.parent
363        samples = payload.get("samples") or []
364        if not isinstance(samples, list):
365            raise ManifestError(
366                f"Manifest 'samples' must be a list, got {type(samples).__name__}",
367            )
368        if max_samples is not None:
369            if max_samples < 1:
370                raise ConfigError(
371                    f"max_samples must be >= 1, got {max_samples}",
372                )
373            samples = samples[:max_samples]
374        return cls(
375            samples=samples,
376            task=TaskType.VIDEO_DEPTH,
377            root=root,
378            batch_size=batch_size,
379            frame_budget=frame_budget,
380            sampling=sampling,
381            compute_fscore=compute_fscore,
382        )
383
384
385# --------------------------------------------------------------------------- #
386# Frame-selection helpers
387# --------------------------------------------------------------------------- #
388
389
390def _stride_sample(T_full: int, budget: int) -> np.ndarray:
391    """Uniform-stride frame selection with first/last anchored.
392
393    Equivalent to FPS in 1D timestamp space at constant frame rate.
394    Always includes indices 0 and ``T_full - 1`` so the clip endpoints
395    are preserved (matters for temporal-edge metrics).
396    """
397    if budget == T_full:
398        return np.arange(T_full, dtype=np.int32)
399    # np.linspace(0, T_full-1, budget) gives evenly-spaced floats with
400    # both endpoints; round and dedupe.
401    raw = np.linspace(0, T_full - 1, budget)
402    idx = np.unique(np.round(raw).astype(np.int32))
403    # In rare cases rounding collapses adjacent picks; pad with the
404    # nearest missing indices to hit the exact budget.
405    if len(idx) < budget:
406        all_idx = set(idx.tolist())
407        for cand in range(T_full):
408            if cand not in all_idx:
409                all_idx.add(cand)
410                if len(all_idx) == budget:
411                    break
412        idx = np.sort(np.fromiter(all_idx, dtype=np.int32))
413    return idx[:budget]
414
415
416def _fps_se3(positions: np.ndarray, budget: int) -> np.ndarray:
417    """Farthest-point sampling in 3D camera-translation space.
418
419    Anchors the first frame (index 0) and iteratively adds the index
420    whose minimum Euclidean distance to the selected set is largest.
421    Deterministic given ``(positions, budget)``.
422
423    Parameters
424    ----------
425    positions : (T_full, 3) float64
426        T265 camera positions in metres.
427    budget : int
428        Number of frames to select. Must satisfy ``2 <= budget <= T_full``.
429
430    Returns
431    -------
432    np.ndarray
433        Sorted indices, shape ``(budget,)``, dtype int32.
434    """
435    if positions.ndim != 2 or positions.shape[1] != 3:
436        raise ConfigError(
437            f"_fps_se3: expected (T, 3) positions, got {positions.shape}",
438            hint="Pass T265 translations as an (N, 3) float64 array.",
439        )
440    T_full = positions.shape[0]
441    if not (2 <= budget <= T_full):
442        raise ConfigError(
443            f"_fps_se3: budget {budget} out of range [2, {T_full}]",
444            hint="Budget must be in [2, len(positions)].",
445        )
446
447    selected = [0]
448    # Squared distance from each frame to the closest selected frame.
449    sq_dist_to_set = np.sum((positions - positions[0]) ** 2, axis=1)
450    sq_dist_to_set[0] = -np.inf  # never re-pick
451
452    while len(selected) < budget:
453        next_idx = int(np.argmax(sq_dist_to_set))
454        selected.append(next_idx)
455        # Update min-distance bookkeeping.
456        new_sq = np.sum((positions - positions[next_idx]) ** 2, axis=1)
457        sq_dist_to_set = np.minimum(sq_dist_to_set, new_sq)
458        sq_dist_to_set[next_idx] = -np.inf
459
460    return np.sort(np.array(selected, dtype=np.int32))
461
462
463def _se3_from_pos_quat(position: np.ndarray, quat_xyzw: np.ndarray) -> np.ndarray:
464    """Pack a (3,) position + (4,) xyzw quaternion into a 4×4 SE(3)."""
465    qx, qy, qz, qw = quat_xyzw
466    # Standard xyzw → rotation matrix
467    R = np.array(
468        [
469            [1 - 2 * (qy * qy + qz * qz), 2 * (qx * qy - qz * qw), 2 * (qx * qz + qy * qw)],
470            [2 * (qx * qy + qz * qw), 1 - 2 * (qx * qx + qz * qz), 2 * (qy * qz - qx * qw)],
471            [2 * (qx * qz - qy * qw), 2 * (qy * qz + qx * qw), 1 - 2 * (qx * qx + qy * qy)],
472        ],
473        dtype=np.float64,
474    )
475    T_mat = np.eye(4, dtype=np.float64)
476    T_mat[:3, :3] = R
477    T_mat[:3, 3] = position
478    return T_mat
479
480
481def _read_translation_track(root: Path, pose_files: List[str]) -> np.ndarray:
482    """Read pose files and return ``(T, 3) float64`` translation track.
483
484    Uses :func:`~rpx_benchmark.decode_contracts.safe_load_pose` to honour
485    both ``.npz`` (legacy) and ``.npy`` (v2) on-disk formats.
486    """
487    positions = []
488    for p in pose_files:
489        pos, _quat = safe_load_pose(root / p)
490        positions.append(pos)
491    return np.stack(positions).astype(np.float64)
492
493
494def _swap_modality(rel_path: str, src: str, dst: str) -> str:
495    """``scenes/x/0/rgb/00000.png`` → ``scenes/x/0/depth/00000.png``.
496
497    Only swaps the *first* occurrence of ``/src/`` to avoid corrupting
498    paths that happen to mention ``rgb`` or ``depth`` elsewhere.
499    """
500    needle = f"/{src}/"
501    repl = f"/{dst}/"
502    i = rel_path.find(needle)
503    if i < 0:
504        return rel_path  # caller will surface a missing-file error downstream
505    return rel_path[:i] + repl + rel_path[i + len(needle) :]
506
507
508# --------------------------------------------------------------------------- #
509# Per-clip scale-and-shift alignment for affine-invariant video models
510# --------------------------------------------------------------------------- #
511
512
513def align_scale_and_shift_per_clip(
514    pred_seq: np.ndarray,
515    gt_seq: np.ndarray,
516    valid_mask_seq: np.ndarray,
517) -> np.ndarray:
518    """Per-clip scale + shift alignment — thin wrapper around the shared
519    pooled solver.
520
521    Delegates to :func:`~rpx_benchmark.metrics.depth_alignment.align_pred_to_gt_pooled`
522    in ``ls_affine`` mode so the math lives in exactly one place
523    (shared with the Image Depth per-scene-phase aligner). Per-clip (not
524    per-frame) alignment is the standard convention for video depth —
525    per-frame alignment would zero out the temporal-consistency
526    signal the video metrics measure.
527
528    The returned sequence is ``s * pred_seq + t`` so callers can plug
529    it straight into the metric calculators.
530    """
531    from .metrics.depth_alignment import align_pred_to_gt_pooled
532
533    if pred_seq.shape != gt_seq.shape or pred_seq.shape != valid_mask_seq.shape:
534        raise ConfigError(
535            f"align_scale_and_shift_per_clip: shape mismatch "
536            f"pred={pred_seq.shape} gt={gt_seq.shape} valid={valid_mask_seq.shape}",
537            hint="All three sequences must be (T, H, W) with matching T/H/W.",
538        )
539    return align_pred_to_gt_pooled(
540        pred_seq=pred_seq,
541        gt_seq=gt_seq,
542        mode="ls_affine",
543        valid_seq=valid_mask_seq.astype(bool),
544    )
545
546
547# Backwards-compatibility alias. Earlier branches (jishnu/d1v-runner-
548# first-adapter, jishnu/d1v-manifest-root-warning) and external callers
549# import ``D1VDataset`` directly. Keeping the alias avoids forcing every
550# caller to update their import line in lockstep with the rename.
551D1VDataset = VideoDepthDataset
552
553
554__all__ = [
555    "VideoDepthDataset",
556    "D1VDataset",
557    "_fps_se3",
558    "_stride_sample",
559    "align_scale_and_shift_per_clip",
560]