rpx_benchmark.data.clip_dataset

  1"""Per-(scene, phase) clip iteration for Video Depth and other
  2sequence-level tasks (e.g. D3 tracking, which is also clip-based).
  3
  4The flat :class:`~rpx_benchmark.loader.RPXDataset` yields one independent
  5:class:`~rpx_benchmark.api.Sample` per frame — correct for per-frame
  6tasks (Image Depth, D2, D4/D5).  Video and tracking tasks instead need the
  7whole phase clip in one shot.  :class:`PhaseClipDataset` groups a depth
  8manifest's frames by ``(scene, phase)``, orders them by frame index, and
  9yields a :class:`PhaseClip` per group, decoding every modality through the
 10*existing* ``RPXDataset`` loaders — no duplicate IO code, one source of
 11truth for how RGB / depth / pose are read.
 12
 13Design notes (``benchmark/README.md`` §3):
 14
 15* Iteration order is ``outer = scene (sorted), inner = phase (0, 1, 2)``,
 16  which matches the cell-log grouping key ``(model, task, scene, phase)``.
 17* Clips are loaded lazily, one at a time, so host-RAM stays at ~1 clip
 18  (~0.6 GB) regardless of dataset size.
 19* Clips shorter than :data:`~rpx_benchmark.metrics.depth_temporal.MIN_CLIP_FRAMES`
 20  are **skipped** (logged) — temporal metrics are meaningless on a few
 21  frames, and we never pad (padding biases OPW/TAE downward).
 22"""
 23
 24from __future__ import annotations
 25
 26import re
 27from dataclasses import dataclass
 28from pathlib import Path
 29from typing import Any, Dict, Iterator, List, Optional, Tuple
 30
 31import numpy as np
 32
 33from ..api import Difficulty, Phase
 34from ..loader import RPXDataset
 35from ..logging_utils import get_logger
 36from ..metrics.depth_alignment import DEPTH_MAX_M, DEPTH_MIN_M
 37from ..metrics.depth_temporal import MIN_CLIP_FRAMES
 38
 39log = get_logger(__name__)
 40
 41__all__ = ["PhaseClip", "PhaseClipDataset", "D1VClip", "D1VClipDataset"]
 42
 43#: Canonical phase ordering for clip iteration and the temporal axis.
 44_PHASE_TO_IDX: Dict[Phase, int] = {Phase.CLUTTER: 0, Phase.INTERACTION: 1, Phase.CLEAN: 2}
 45
 46#: Default D435 intrinsics (shared with the metric calculators).
 47_D435_K = np.array([[615.0, 0.0, 320.0], [0.0, 615.0, 240.0], [0.0, 0.0, 1.0]], dtype=np.float64)
 48
 49_TRAILING_INT = re.compile(r"(\d+)(?!.*\d)")
 50
 51
 52@dataclass
 53class PhaseClip:
 54    """One phase clip for Video Depth evaluation (decisions doc §3.2).
 55
 56    Attributes
 57    ----------
 58    rgb_seq : (T, H, W, 3) uint8
 59    depth_gt_seq : (T, H, W) float32, metres (0 = invalid sentinel).
 60    valid_mask_seq : (T, H, W) bool — GT in (0.3, 5.0) m and finite.
 61    scene_id : str
 62    phase_idx : int — 0=clutter, 1=interaction, 2=clean.
 63    frame_indices : (T,) int — original frame numbers, ascending.
 64    intrinsics : (3, 3) float64 — pinhole K (default D435).
 65    poses : (T, 4, 4) float64 | None — world-from-camera SE(3); ``None``
 66        if any frame in the clip lacks a T265 pose.
 67    difficulty : Difficulty | None — the scene's ESD tier (easy/medium/
 68        hard), needed so the (scene, phase) cell can carry its split tag
 69        for cross-split (phase×difficulty) Φ. ``None`` if untagged.
 70    """
 71
 72    rgb_seq: np.ndarray
 73    depth_gt_seq: np.ndarray
 74    valid_mask_seq: np.ndarray
 75    scene_id: str
 76    phase_idx: int
 77    frame_indices: np.ndarray
 78    intrinsics: np.ndarray
 79    poses: Optional[np.ndarray]
 80    difficulty: Optional[Difficulty] = None
 81
 82    @property
 83    def num_frames(self) -> int:
 84        return int(self.rgb_seq.shape[0])
 85
 86    @property
 87    def phase(self) -> Phase:
 88        return {0: Phase.CLUTTER, 1: Phase.INTERACTION, 2: Phase.CLEAN}[self.phase_idx]
 89
 90
 91def _phase_idx_of(entry: Dict[str, Any]) -> Optional[int]:
 92    """Resolve an entry's phase to a canonical 0/1/2 index, or ``None``."""
 93    v = entry.get("phase")
 94    if v is None:
 95        return None
 96    if isinstance(v, (int, np.integer)):
 97        return int(v) if int(v) in (0, 1, 2) else None
 98    try:
 99        return _PHASE_TO_IDX[Phase(v)]
100    except (ValueError, KeyError):
101        return None
102
103
104def _frame_key(entry: Dict[str, Any]) -> int:
105    """Best-effort ascending frame index for ordering within a clip.
106
107    Prefers an explicit ``frame`` field (top-level or in ``metadata``),
108    then the trailing integer of the sample ``id``, then the RGB filename
109    stem.  Falls back to 0 so ordering is at worst stable-by-insertion.
110    """
111    for src in (entry.get("frame"), (entry.get("metadata") or {}).get("frame")):
112        if isinstance(src, (int, np.integer)):
113            return int(src)
114    for text in (str(entry.get("id", "")), Path(str(entry.get("rgb", ""))).stem):
115        m = _TRAILING_INT.search(text)
116        if m:
117            return int(m.group(1))
118    return 0
119
120
121@dataclass
122class PhaseClipDataset:
123    """Groups a depth manifest into per-(scene, phase) clips.
124
125    Construct via :meth:`from_manifest`.  Iterating yields
126    :class:`PhaseClip` objects in ``(scene, phase)`` order, skipping clips
127    with fewer than ``min_frames`` frames.
128    """
129
130    base: RPXDataset
131    min_frames: int = MIN_CLIP_FRAMES
132    intrinsics: np.ndarray = None  # type: ignore[assignment]
133
134    def __post_init__(self) -> None:
135        if self.intrinsics is None:
136            self.intrinsics = _D435_K
137
138    @classmethod
139    def from_manifest(
140        cls,
141        manifest_path: str | Path,
142        min_frames: int = MIN_CLIP_FRAMES,
143        intrinsics: Optional[np.ndarray] = None,
144    ) -> "PhaseClipDataset":
145        base = RPXDataset.from_manifest(manifest_path, batch_size=1)
146        return cls(base=base, min_frames=min_frames,
147                   intrinsics=intrinsics if intrinsics is not None else _D435_K)
148
149    # -- grouping ---------------------------------------------------------- #
150
151    def _groups(self) -> List[Tuple[str, int, List[Dict[str, Any]]]]:
152        """Return ``[(scene, phase_idx, [entries...]), ...]`` in iteration order.
153
154        Entries are sorted by frame index within each group; groups are
155        sorted by ``(scene, phase_idx)``.  Entries with no resolvable
156        scene or phase are dropped (warned once per dropped count).
157        """
158        buckets: Dict[Tuple[str, int], List[Dict[str, Any]]] = {}
159        dropped = 0
160        for entry in self.base.samples:
161            scene = entry.get("scene") or (entry.get("metadata") or {}).get("scene_id")
162            phase_idx = _phase_idx_of(entry)
163            if scene is None or phase_idx is None:
164                dropped += 1
165                continue
166            buckets.setdefault((str(scene), phase_idx), []).append(entry)
167        if dropped:
168            log.warning("PhaseClipDataset: dropped %d entries with no scene/phase", dropped)
169
170        groups = []
171        for (scene, phase_idx), entries in sorted(buckets.items(), key=lambda kv: kv[0]):
172            entries_sorted = sorted(entries, key=_frame_key)
173            groups.append((scene, phase_idx, entries_sorted))
174        return groups
175
176    def clip_plan(self) -> List[Tuple[str, int, int]]:
177        """``[(scene, phase_idx, n_frames), ...]`` without loading pixels.
178
179        Cheap introspection for run planning / progress bars; reflects the
180        ``min_frames`` skip.
181        """
182        return [
183            (scene, phase_idx, len(entries))
184            for scene, phase_idx, entries in self._groups()
185            if len(entries) >= self.min_frames
186        ]
187
188    # -- loading ----------------------------------------------------------- #
189
190    def _load_clip(self, scene: str, phase_idx: int, entries: List[Dict[str, Any]]) -> PhaseClip:
191        rgbs, depths, valids, frame_ids, poses = [], [], [], [], []
192        pose_complete = True
193        difficulty: Optional[Difficulty] = None
194        for entry in entries:
195            sample = self.base._load_sample(entry)
196            depth = np.asarray(sample.ground_truth.depth_map, dtype=np.float32)
197            rgbs.append(np.asarray(sample.rgb, dtype=np.uint8))
198            depths.append(depth)
199            valids.append(np.isfinite(depth) & (depth > DEPTH_MIN_M) & (depth < DEPTH_MAX_M))
200            frame_ids.append(_frame_key(entry))
201            if sample.camera_pose is None:
202                pose_complete = False
203            else:
204                poses.append(np.asarray(sample.camera_pose, dtype=np.float64))
205            # Difficulty is a per-scene ESD tier — constant within a
206            # (scene, phase) clip; take the first non-null and warn on drift.
207            if sample.difficulty is not None:
208                if difficulty is None:
209                    difficulty = sample.difficulty
210                elif difficulty != sample.difficulty:
211                    log.warning(
212                        "PhaseClipDataset: %s phase %d has mixed difficulty (%s vs %s); keeping %s",
213                        scene, phase_idx, difficulty, sample.difficulty, difficulty,
214                    )
215
216        return PhaseClip(
217            rgb_seq=np.stack(rgbs),
218            depth_gt_seq=np.stack(depths),
219            valid_mask_seq=np.stack(valids),
220            scene_id=scene,
221            phase_idx=phase_idx,
222            frame_indices=np.asarray(frame_ids, dtype=np.int64),
223            intrinsics=self.intrinsics,
224            poses=np.stack(poses) if (pose_complete and poses) else None,
225            difficulty=difficulty,
226        )
227
228    def __iter__(self) -> Iterator[PhaseClip]:
229        for scene, phase_idx, entries in self._groups():
230            if len(entries) < self.min_frames:
231                log.warning(
232                    "PhaseClipDataset: skipping %s phase %d — only %d frames (< %d)",
233                    scene, phase_idx, len(entries), self.min_frames,
234                )
235                continue
236            yield self._load_clip(scene, phase_idx, entries)
237
238    def __len__(self) -> int:
239        """Number of clips that will be yielded (after the min_frames skip)."""
240        return len(self.clip_plan())
241
242
243# --------------------------------------------------------------------------- #
244# Paper-facing aliases — RPX depth task labels onto the generic clip API.
245# The generic names above are canonical; these are the paper's Video Depth labels.
246# --------------------------------------------------------------------------- #
247
248#: Paper alias for :class:`PhaseClip` (RPX task Video Depth).
249D1VClip = PhaseClip
250#: Paper alias for :class:`PhaseClipDataset` (RPX task Video Depth).
251D1VClipDataset = PhaseClipDataset