rpx_benchmark.dataset_hub.recipes

  1"""Task → modality recipes for selective dataset downloads.
  2
  3A *recipe* answers the question "what subset of files does the dataset need
  4to materialise on disk so a given task can run?". The downloader takes the
  5recipe, expands it to a list of glob patterns, and hands them to
  6``huggingface_hub.snapshot_download(allow_patterns=...)``.
  7
  8Two scene families have separate recipe tables because they have different
  9on-disk shapes:
 10
 11* ``MULTI_OBJECT_TASK_RECIPES`` — scenes named ``scene<N>.<building>.<area>/``
 12  with three phases ``0/`` ``1/`` ``2/``. Used for the bulk of perception
 13  benchmarks (segmentation, depth, pose, tracking, ...).
 14
 15* ``SINGLE_OBJECT_TASK_RECIPES`` — scenes named ``object<N>.<name>/`` with
 16  one collection ``0/`` (a 360° turntable view). Used for in-context tasks
 17  (object templates, few-shot reference views, retrieval).
 18"""
 19
 20from __future__ import annotations
 21
 22from dataclasses import dataclass, field
 23from enum import Enum
 24from typing import Dict, FrozenSet, Iterable, Optional
 25
 26from ..exceptions import ConfigError
 27
 28DEFAULT_REPO_ID = "IRVLUTD/RPX"
 29
 30
 31class SceneType(str, Enum):
 32    """Top-level partition: which scene family a recipe targets."""
 33
 34    MULTI_OBJECT = "multi_object"
 35    SINGLE_OBJECT = "single_object"
 36    EGO = "ego"
 37
 38
 39@dataclass(frozen=True)
 40class TaskRecipe:
 41    """Declarative spec: 'task X needs these modalities for inputs and labels'.
 42
 43    The split between ``inputs`` and ``labels`` is purely a documentation
 44    aid — the downloader unions them when expanding to glob patterns.
 45    Storing them separately makes it obvious *why* each modality is in the
 46    recipe (which matters when modalities are added or renamed).
 47    """
 48
 49    name: str
 50    scene_type: SceneType
 51    inputs: FrozenSet[str]
 52    labels: FrozenSet[str] = field(default_factory=frozenset)
 53    notes: str = ""
 54
 55    def all_modalities(self) -> FrozenSet[str]:
 56        return self.inputs | self.labels
 57
 58
 59# --------------------------------------------------------------------- #
 60# Modality vocabulary (matches the on-disk capture layout under
 61# `test_dataset_aggregated/<scene>/<phase>/<modality>/`)
 62# --------------------------------------------------------------------- #
 63
 64#: Raw modalities — captured once, never re-released.
 65RGB = "rgb"
 66DEPTH = "depth"
 67FISHEYE = "fisheye"
 68CAM_POSE = "cam_pose"
 69
 70#: Label modalities — versioned (``labels/<name>/v1``, ``v2``, ...).
 71MASKS = "masks"  # essential SAM2 output (object masks)
 72MASKS_AUX = "masks_aux"  # opt-in viz/derived (bbox_overlay, dino_output, palette, ...)
 73QUESTIONNAIRE = "questionnaire"  # FewSOL-style per-scene/object text
 74VQA = "vqa"  # generated Q&A pairs (lands in v1.x)
 75
 76
 77# --------------------------------------------------------------------- #
 78# Recipe tables
 79# --------------------------------------------------------------------- #
 80
 81# Multi-object scenes (3 phases each, ESD-stratified).
 82MULTI_OBJECT_TASK_RECIPES: Dict[str, TaskRecipe] = {
 83    "monocular_depth": TaskRecipe(
 84        name="monocular_depth",
 85        scene_type=SceneType.MULTI_OBJECT,
 86        inputs=frozenset({RGB}),
 87        labels=frozenset({DEPTH}),
 88        notes=(
 89            "Paper task Image Depth. Frame-level: model receives one RGB frame, "
 90            "outputs one depth map. Depth captured by D435 is the GT."
 91        ),
 92    ),
 93    "video_depth": TaskRecipe(
 94        name="video_depth",
 95        scene_type=SceneType.MULTI_OBJECT,
 96        # Same modalities as monocular_depth — Video Depth differs only in
 97        # iteration unit (per-(scene, phase) clip) and metric set
 98        # (adds temporal metrics on top of Image Depth's per-frame ones).
 99        # Downloading via `download_for_task("video_depth", ...)` pulls
100        # the exact same tar shards as monocular_depth; the runner
101        # decides how to feed them to the model.
102        inputs=frozenset({RGB}),
103        labels=frozenset({DEPTH}),
104        notes=(
105            "Paper task Video Depth. Video-level: model receives the full "
106            "~250-frame phase clip as RGB sequence, outputs a per-frame "
107            "depth sequence. Same scenes and GT as monocular_depth; "
108            "differs in iteration unit and metric set (per-frame metrics "
109            "+ temporal metrics — TODO: finalise from Feynman's research "
110            "in benchmark/README.md)."
111        ),
112    ),
113    "rgbd_segmentation": TaskRecipe(
114        name="rgbd_segmentation",
115        scene_type=SceneType.MULTI_OBJECT,
116        inputs=frozenset({RGB, DEPTH}),
117        labels=frozenset({MASKS}),
118    ),
119    "segmentation": TaskRecipe(
120        name="segmentation",
121        scene_type=SceneType.MULTI_OBJECT,
122        inputs=frozenset({RGB}),
123        labels=frozenset({MASKS}),
124    ),
125    "relative_pose": TaskRecipe(
126        name="relative_pose",
127        scene_type=SceneType.MULTI_OBJECT,
128        inputs=frozenset({RGB}),
129        labels=frozenset({CAM_POSE}),
130    ),
131    "rgbd_relative_pose": TaskRecipe(
132        name="rgbd_relative_pose",
133        scene_type=SceneType.MULTI_OBJECT,
134        inputs=frozenset({RGB, DEPTH}),
135        labels=frozenset({CAM_POSE}),
136    ),
137    "stereo_depth": TaskRecipe(
138        name="stereo_depth",
139        scene_type=SceneType.MULTI_OBJECT,
140        inputs=frozenset({FISHEYE}),
141        labels=frozenset({DEPTH}),
142        notes="T265 stereo fisheye → depth. D435 depth is the GT after rectification.",
143    ),
144    "object_tracking": TaskRecipe(
145        name="object_tracking",
146        scene_type=SceneType.MULTI_OBJECT,
147        inputs=frozenset({RGB}),
148        labels=frozenset({MASKS}),
149    ),
150    "vqa": TaskRecipe(
151        name="vqa",
152        scene_type=SceneType.MULTI_OBJECT,
153        inputs=frozenset({RGB}),
154        labels=frozenset({VQA, QUESTIONNAIRE}),
155        notes="Reserved slot — VQA labels land in v1.x.",
156    ),
157}
158
159# Ego (egocentric/GoPro) scenes — one continuous capture per scene, no
160# phase subdivision (mirrors the single-object shape: one collection "0").
161# Same scene_id as the sibling MOS scene (e.g. "scene20.su.checkerboard"),
162# so downloaders can pair an ego clip with its static-camera MOS phases.
163# Only rgb + masks exist for ego today — no depth/fisheye/cam_pose.
164EGO_TASK_RECIPES: Dict[str, TaskRecipe] = {
165    "ego_segmentation": TaskRecipe(
166        name="ego_segmentation",
167        scene_type=SceneType.EGO,
168        inputs=frozenset({RGB}),
169        labels=frozenset({MASKS}),
170        notes="Frame-level instance segmentation from the egocentric (head/hand-worn) viewpoint.",
171    ),
172    "ego_object_tracking": TaskRecipe(
173        name="ego_object_tracking",
174        scene_type=SceneType.EGO,
175        inputs=frozenset({RGB}),
176        labels=frozenset({MASKS}),
177        notes=(
178            "Same modalities as ego_segmentation — separate recipe because the "
179            "task (per-frame masks vs. track-consistent masks across the clip) "
180            "differs, matching the mos segmentation/object_tracking split."
181        ),
182    ),
183    "ego_vqa": TaskRecipe(
184        name="ego_vqa",
185        scene_type=SceneType.EGO,
186        inputs=frozenset({RGB}),
187        labels=frozenset({VQA, QUESTIONNAIRE}),
188        notes="Reserved slot, mirrors mos vqa — VQA labels land in v1.x.",
189    ),
190}
191
192# Single-object scenes (1 collection, used for in-context / template tasks).
193SINGLE_OBJECT_TASK_RECIPES: Dict[str, TaskRecipe] = {
194    "object_templates": TaskRecipe(
195        name="object_templates",
196        scene_type=SceneType.SINGLE_OBJECT,
197        inputs=frozenset({RGB}),
198        labels=frozenset({MASKS}),
199        notes="360° turntable RGB + mask — feed as reference views for in-context tasks.",
200    ),
201    "object_templates_rgbd": TaskRecipe(
202        name="object_templates_rgbd",
203        scene_type=SceneType.SINGLE_OBJECT,
204        inputs=frozenset({RGB, DEPTH}),
205        labels=frozenset({MASKS}),
206    ),
207    "object_pose_library": TaskRecipe(
208        name="object_pose_library",
209        scene_type=SceneType.SINGLE_OBJECT,
210        inputs=frozenset({RGB, DEPTH}),
211        labels=frozenset({CAM_POSE, MASKS}),
212        notes="Per-object pose graph for retrieval / pose initialisation.",
213    ),
214}
215
216
217def resolve_recipe(
218    task: str,
219    scene_type: Optional[SceneType] = None,
220) -> TaskRecipe:
221    """Look up a recipe by ``task`` name.
222
223    If ``scene_type`` is given, only that table is searched (use this when a
224    name collides between multi- and single-object tables). Otherwise we
225    search multi-object first, then single-object, and raise on miss.
226    """
227    if scene_type is SceneType.MULTI_OBJECT:
228        return MULTI_OBJECT_TASK_RECIPES[task]
229    if scene_type is SceneType.SINGLE_OBJECT:
230        return SINGLE_OBJECT_TASK_RECIPES[task]
231    if scene_type is SceneType.EGO:
232        return EGO_TASK_RECIPES[task]
233    if task in MULTI_OBJECT_TASK_RECIPES:
234        return MULTI_OBJECT_TASK_RECIPES[task]
235    if task in SINGLE_OBJECT_TASK_RECIPES:
236        return SINGLE_OBJECT_TASK_RECIPES[task]
237    if task in EGO_TASK_RECIPES:
238        return EGO_TASK_RECIPES[task]
239    raise ConfigError(
240        f"Unknown task {task!r}.",
241        hint=(
242            "Known multi-object recipes: "
243            f"{sorted(MULTI_OBJECT_TASK_RECIPES)}. "
244            "Known single-object recipes: "
245            f"{sorted(SINGLE_OBJECT_TASK_RECIPES)}. "
246            "Known ego recipes: "
247            f"{sorted(EGO_TASK_RECIPES)}. "
248            "Pass `scene_type=SceneType.SINGLE_OBJECT` if a name collides "
249            "between tables."
250        ),
251    )
252
253
254def all_recipe_names(scene_type: Optional[SceneType] = None) -> Iterable[str]:
255    """Iterator over recipe names, optionally filtered by scene type."""
256    if scene_type is None or scene_type is SceneType.MULTI_OBJECT:
257        yield from MULTI_OBJECT_TASK_RECIPES
258    if scene_type is None or scene_type is SceneType.SINGLE_OBJECT:
259        yield from SINGLE_OBJECT_TASK_RECIPES
260    if scene_type is None or scene_type is SceneType.EGO:
261        yield from EGO_TASK_RECIPES