rpx_benchmark.dataset_hub.split_manifests

  1"""Per-task, per-split manifest writer for the dataset hub.
  2
  3The toolkit's :func:`rpx_benchmark.hub.download_split` expects a JSON file at
  4``manifests/<recipe_key>/<split>.json`` on the published HF repo. The hub's
  5existing ``manifest.py`` only emits the all-frames Parquet
  6(``manifest/frames_v1.parquet``); this module fills the gap so a freshly
  7published HF release works for every task the loader supports.
  8
  9Per-task entry contracts
 10------------------------
 11Each TaskType the loader can parse expects a specific set of keys on each
 12manifest entry. Using the wrong key (``cam_pose`` instead of ``pose``,
 13``masks`` instead of ``mask``, ``fisheye`` instead of
 14``fisheye_left``/``fisheye_right``) is a silent 404 at load time.
 15
 16This module owns the per-task contract:
 17
 18==================  ===================  ============================================
 19Recipe key          TaskType.value       Required entry keys + GT keys
 20==================  ===================  ============================================
 21monocular_depth     monocular_depth      rgb, depth
 22segmentation        object_segmentation  rgb, mask
 23rgbd_segmentation   object_segmentation  rgb, depth, mask
 24stereo_depth        monocular_depth      fisheye_left, fisheye_right, depth
 25relative_pose       relative_camera_pose rgb, rgb_b, pose_a, pose_b           [paired]
 26rgbd_relative_pose  relative_camera_pose rgb, depth, rgb_b, depth_b,
 27                                          pose_a, pose_b                       [paired]
 28object_tracking     object_tracking      rgb, mask                            [per-frame]
 29vqa                 visual_grounding     rgb, text                            [needs labels]
 30==================  ===================  ============================================
 31
 32Modality file layout (where the actual files land after extraction)
 33-------------------------------------------------------------------
 34On-disk:
 35
 36==========  ============================================  ============
 37Modality    Subdir under ``extracted/scenes/<s>/<p>/``    Suffix
 38==========  ============================================  ============
 39rgb         ``rgb/``                                       ``.png``
 40depth       ``depth/``                                     ``.png`` (16-bit mm)
 41fisheye     ``fisheye/{left,right}/``                      ``.png``
 42masks       ``sam2/masks/``                                ``.png``
 43cam_pose    ``cam_pose/``                                  ``.npz``
 44==========  ============================================  ============
 45
 46Splits — scene-wise, mode-aggregated
 47------------------------------------
 48Splits are scene-wise (every phase of a scene shares one tier) so STR /
 49cross-phase comparisons aren't biased by tier reassignment.
 50
 51Usage
 52-----
 53::
 54
 55    from rpx_benchmark.dataset_hub.split_manifests import write_split_manifests
 56    written = write_split_manifests("/path/to/staging")
 57    # → {(recipe_key, split): manifest_path, ...}
 58"""
 59
 60from __future__ import annotations
 61
 62import json
 63import logging
 64from pathlib import Path
 65from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple
 66
 67from .packer import SCENE_ROOT_BY_TYPE, phase_segment
 68from .recipes import SceneType
 69
 70log = logging.getLogger(__name__)
 71
 72
 73# ---------------------------------------------------------------------- #
 74# Per-modality on-disk layout (after extraction)
 75# ---------------------------------------------------------------------- #
 76#
 77# (parquet column suffix, on-disk subdir, file extension)
 78#
 79# The parquet's ``has_<m>`` / ``shard_<m>`` columns use the column suffix.
 80# The on-disk subdir is what ``hub._extract_snapshot_tars`` materialises
 81# under ``extracted/scenes/<scene>/<phase>/``.  The file extension is
 82# what each frame ends in once extracted (rgb/depth → .png from the
 83# 16-bit mm, cam_pose → .npz containing the SE(3) matrix).
 84
 85_MODALITY_LAYOUT: Dict[str, Tuple[str, str, str]] = {
 86    # modality_name → (parquet_col_suffix, extracted_subdir, default_ext)
 87    #
 88    # ``default_ext`` is the v1-style extension. Real-world datasets are
 89    # described by ``current.json``'s ``modality_extensions`` block
 90    # (written by the manifest builder); when a per-modality entry is
 91    # present there, it overrides this default at write-split-manifests
 92    # time. Anything missing from ``current.json`` keeps the v1 default
 93    # for backward compatibility with already-uploaded datasets.
 94    "rgb": ("rgb", "rgb", ".png"),
 95    "depth": ("depth", "depth", ".png"),
 96    "fisheye": ("fisheye", "fisheye", ".png"),  # see note for left/right
 97    "fisheye_left": ("fisheye", "fisheye/left", ".png"),
 98    "fisheye_right": ("fisheye", "fisheye/right", ".png"),
 99    "masks": ("masks", "sam2/masks", ".png"),
100    "cam_pose": ("cam_pose", "cam_pose", ".npz"),
101}
102
103
104# Active layout for the duration of one ``write_split_manifests`` call.
105# Module-level rather than a thread-local because this whole subsystem
106# is invoked single-threadedly from the CLI. Reset to ``None`` on the
107# way out so a subsequent call without overrides falls back to defaults.
108_active_layout: Optional[Dict[str, Tuple[str, str, str]]] = None
109
110
111def _current_layout() -> Dict[str, Tuple[str, str, str]]:
112    return _active_layout if _active_layout is not None else _MODALITY_LAYOUT
113
114
115def _layout_with_overrides(
116    extensions: Mapping[str, str],
117) -> Dict[str, Tuple[str, str, str]]:
118    """Return a copy of ``_MODALITY_LAYOUT`` with per-modality extensions
119    swapped in from ``extensions``. Unknown modality names in
120    ``extensions`` are ignored (forward compatibility with any new
121    modality the manifest writer might add later)."""
122    out: Dict[str, Tuple[str, str, str]] = {}
123    for modality, (col, subdir, default_ext) in _MODALITY_LAYOUT.items():
124        ext = extensions.get(modality, default_ext)
125        out[modality] = (col, subdir, ext)
126    return out
127
128
129def _modality_path(
130    scene: str, phase: int, modality: str, frame_stem: str,
131    scene_type: str = SceneType.MULTI_OBJECT.value,
132) -> str:
133    """Build the manifest's path string for one (scene, phase, modality, frame).
134
135    Path is relative to the manifest's ``root`` (set by ``download_split``
136    to the snapshot root). After ``_extract_snapshot_tars`` runs, the file
137    lives under ``extracted/<repo_root>/...`` where ``<repo_root>`` is
138    ``scenes`` for mos, ``objects`` for sos, ``ego`` for ego (see
139    packer.SCENE_ROOT_BY_TYPE — the single source of truth this must stay
140    in sync with). ``scene_type`` defaults to multi_object/"scenes" so
141    every pre-ego call site (all mos-only specs) is unaffected; only the
142    specs actually shared with ego (_SegmentationSpec, _ObjectTrackingSpec)
143    pass the row's real scene_type.
144
145    Reads the active layout (set by :func:`write_split_manifests`); falls
146    back to :data:`_MODALITY_LAYOUT` defaults outside that context.
147    """
148    layout = _current_layout()
149    if modality not in layout:
150        from ..exceptions import ConfigError
151
152        raise ConfigError(
153            f"unknown modality {modality!r}",
154            hint=f"Known modalities: {', '.join(sorted(layout))}",
155        )
156    _, subdir, ext = layout[modality]
157    st = SceneType(scene_type)
158    repo_root = SCENE_ROOT_BY_TYPE[st]
159    seg = phase_segment(st, phase)
160    return f"extracted/{repo_root}/{scene}/{seg}/{subdir}/{frame_stem}{ext}"
161
162
163def _has_col(modality: str) -> str:
164    """Return the parquet ``has_<col>`` column name for a modality."""
165    col_suffix, _, _ = _MODALITY_LAYOUT[modality]
166    return f"has_{col_suffix}"
167
168
169def _load_modality_extensions(staging_root: Path) -> Dict[str, str]:
170    """Read ``modality_extensions`` from ``<staging>/manifest/current.json``.
171
172    Returns an empty dict if the file is missing, malformed, or
173    pre-dates the ``modality_extensions`` field — that's the backward-
174    compatible signal to keep the v1 hardcoded extensions.
175    """
176    cur_path = staging_root / "manifest" / "current.json"
177    if not cur_path.is_file():
178        return {}
179    try:
180        payload = json.loads(cur_path.read_text(encoding="utf-8"))
181    except json.JSONDecodeError:
182        return {}
183    raw = payload.get("modality_extensions") or {}
184    if not isinstance(raw, dict):
185        return {}
186    return {str(k): str(v) for k, v in raw.items()}
187
188
189# ---------------------------------------------------------------------- #
190# Per-task spec — links recipe keys → TaskType value + entry-builder
191# ---------------------------------------------------------------------- #
192
193
194class _TaskSpec:
195    """How to build manifest entries for one recipe key.
196
197    Subclasses override ``build_entries`` to return a list of dicts
198    matching the loader's expected entry shape for that task.
199    """
200
201    #: The string the loader's ``TaskType()`` constructor needs to see.
202    task_type_value: str = ""
203
204    #: Required modalities, by parquet column suffix. Used to filter the
205    #: parquet rows (a row enters the manifest only if every required
206    #: modality is present).
207    required_modalities: Tuple[str, ...] = ()
208
209    #: scan/manifest's SceneType.value strings this spec applies to.
210    #: Defaults to multi_object (mos) — every spec below predates ego and
211    #: was written assuming only mos rows would ever reach it (nothing here
212    #: checked scene_type at all, so ego rows silently satisfied the same
213    #: rgb+masks requirement and leaked into e.g. "segmentation" instead of
214    #: a distinct "ego_segmentation" — see the Ego* specs below, which are
215    #: the fix, not new functionality).
216    scene_types: Tuple[str, ...] = ("multi_object",)
217
218    def build_entries(self, df) -> List[Dict[str, Any]]:
219        raise NotImplementedError
220
221
222class _SingleFrameDepthSpec(_TaskSpec):
223    """``monocular_depth``: rgb + depth GT, one entry per frame."""
224
225    task_type_value = "monocular_depth"
226    required_modalities = ("rgb", "depth")
227
228    def build_entries(self, df) -> List[Dict[str, Any]]:
229        return [_single_frame_entry(row, ["rgb", "depth"]) for _, row in df.iterrows()]
230
231
232class _VideoDepthSpec(_TaskSpec):
233    """``video_depth``: rgb + depth GT, one entry per ``(scene, phase)`` clip.
234
235    Emits ONE entry per (scene, phase) — not per-frame — with the
236    relevant frame filename lists pre-resolved so ``VideoDepthDataset`` can
237    iterate clips directly. The structure differs from every other
238    spec on this page because Video Depth's iteration unit is a clip, not a
239    frame; the cell-log key for the runner becomes
240    ``(model, task, scene, phase, frame_budget)`` and one row per clip
241    is the right granularity.
242    """
243
244    task_type_value = "video_depth"
245    required_modalities = ("rgb", "depth")
246
247    def build_entries(self, df) -> List[Dict[str, Any]]:
248        # Group the per-frame parquet rows by (scene_id, phase) and
249        # collapse each group into a single clip entry.
250        rows: List[Dict[str, Any]] = []
251        # GroupBy preserves first-row metadata for scene_type/difficulty
252        # which are constant within a (scene, phase) cell.
253        for (scene, phase), group in df.groupby(["scene_id", "phase"], sort=True):
254            stems = sorted(_frame_stem(r) for _, r in group.iterrows())
255            rgb_paths = [_modality_path(scene, int(phase), "rgb", s) for s in stems]
256            depth_paths = [_modality_path(scene, int(phase), "depth", s) for s in stems]
257            pose_paths = [_modality_path(scene, int(phase), "cam_pose", s) for s in stems]
258            # Use the first row for cell-level metadata (split, difficulty,
259            # scene_type). All rows in this group carry the same values.
260            first = next(iter(group.itertuples()))
261            rows.append(
262                {
263                    "id": f"{scene}__{int(phase)}",
264                    "scene_id": scene,
265                    "phase": int(phase),
266                    "scene_type": getattr(first, "scene_type", None),
267                    "difficulty": str(getattr(first, "split", "")) or None,
268                    "frame_filenames": rgb_paths,
269                    "depth_filenames": depth_paths,
270                    "pose_filenames": pose_paths,
271                    "metadata": {
272                        "scene_id": scene,
273                        "phase_idx": int(phase),
274                        "n_frames": len(stems),
275                    },
276                }
277            )
278        return rows
279
280
281class _SegmentationSpec(_TaskSpec):
282    """``segmentation``: rgb + mask (singular!) GT.
283
284    Loader's ``_load_ground_truth`` does ``self._load_mask(entry["mask"])``;
285    the entry key is ``mask`` (singular) but the modality directory is
286    ``sam2/masks``. We rename the entry key explicitly here.
287    """
288
289    task_type_value = "object_segmentation"
290    required_modalities = ("rgb", "masks")
291
292    def build_entries(self, df) -> List[Dict[str, Any]]:
293        rows: List[Dict[str, Any]] = []
294        for _, row in df.iterrows():
295            e = _base_entry(row)
296            scene = e["scene_id"]
297            phase = e["phase"]
298            stem = _frame_stem(row)
299            st = str(row["scene_type"])
300            e["rgb"] = _modality_path(scene, phase, "rgb", stem, scene_type=st)
301            e["mask"] = _modality_path(scene, phase, "masks", stem, scene_type=st)
302            rows.append(e)
303        return rows
304
305
306class _RGBDSegmentationSpec(_TaskSpec):
307    """``rgbd_segmentation``: rgb + depth + mask. Loader's task is still
308    OBJECT_SEGMENTATION (depth is an extra input the model sees, not GT)."""
309
310    task_type_value = "object_segmentation"
311    required_modalities = ("rgb", "depth", "masks")
312
313    def build_entries(self, df) -> List[Dict[str, Any]]:
314        rows: List[Dict[str, Any]] = []
315        for _, row in df.iterrows():
316            e = _base_entry(row)
317            scene = e["scene_id"]
318            phase = e["phase"]
319            stem = _frame_stem(row)
320            e["rgb"] = _modality_path(scene, phase, "rgb", stem)
321            e["depth"] = _modality_path(scene, phase, "depth", stem)
322            e["mask"] = _modality_path(scene, phase, "masks", stem)
323            rows.append(e)
324        return rows
325
326
327class _StereoDepthSpec(_TaskSpec):
328    """``stereo_depth``: fisheye_left + fisheye_right inputs, depth GT.
329
330    Loader's TaskType is MONOCULAR_DEPTH (depth GT is the same;
331    stereo-vs-mono is an input-side distinction).
332    """
333
334    task_type_value = "monocular_depth"
335    required_modalities = ("rgb", "depth", "fisheye")
336
337    def build_entries(self, df) -> List[Dict[str, Any]]:
338        rows: List[Dict[str, Any]] = []
339        for _, row in df.iterrows():
340            e = _base_entry(row)
341            scene = e["scene_id"]
342            phase = e["phase"]
343            stem = _frame_stem(row)
344            e["rgb"] = _modality_path(scene, phase, "rgb", stem)
345            e["depth"] = _modality_path(scene, phase, "depth", stem)
346            e["fisheye_left"] = _modality_path(scene, phase, "fisheye_left", stem)
347            e["fisheye_right"] = _modality_path(scene, phase, "fisheye_right", stem)
348            rows.append(e)
349        return rows
350
351
352class _RelativePoseSpec(_TaskSpec):
353    """``relative_pose``: paired-frame task.
354
355    Each entry pairs frame N (`rgb`, `pose_a`) with frame N+stride
356    (`rgb_b`, `pose_b`) within the same scene+phase. Stride defaults to 5
357    frames (~0.5 s at 10 Hz capture; matches existing
358    ``scripts/generate_keypoint_pairs.py`` convention).
359    """
360
361    task_type_value = "relative_camera_pose"
362    required_modalities = ("rgb", "cam_pose")
363    pair_stride: int = 5
364
365    def build_entries(self, df) -> List[Dict[str, Any]]:
366        rows: List[Dict[str, Any]] = []
367        # Group rows by (scene, phase) — pairs only within a phase.
368        group_iter = df.groupby(["scene_id", "phase"], sort=False)
369        for (scene, phase), grp in group_iter:
370            grp_sorted = grp.sort_values("frame_idx").reset_index(drop=True)
371            n = len(grp_sorted)
372            for i in range(n - self.pair_stride):
373                row_a = grp_sorted.iloc[i]
374                row_b = grp_sorted.iloc[i + self.pair_stride]
375                stem_a = _frame_stem(row_a)
376                stem_b = _frame_stem(row_b)
377                e: Dict[str, Any] = {
378                    "id": f"{scene}__{phase}__{stem_a}__{stem_b}",
379                    "scene_id": scene,
380                    "phase": int(phase),
381                    "frame_idx": int(row_a["frame_idx"]),
382                    "frame_idx_b": int(row_b["frame_idx"]),
383                    "difficulty": str(row_a["split"]),
384                    "metadata": {
385                        "scene_id": scene,
386                        "phase_idx": int(phase),
387                        "frame": stem_a,
388                        "frame_b": stem_b,
389                        "pair_stride": self.pair_stride,
390                    },
391                    "rgb": _modality_path(scene, int(phase), "rgb", stem_a),
392                    "rgb_b": _modality_path(scene, int(phase), "rgb", stem_b),
393                    "pose_a": _modality_path(scene, int(phase), "cam_pose", stem_a),
394                    "pose_b": _modality_path(scene, int(phase), "cam_pose", stem_b),
395                }
396                rows.append(e)
397        return rows
398
399
400class _RGBDRelativePoseSpec(_RelativePoseSpec):
401    """``rgbd_relative_pose``: paired with depth on both frames."""
402
403    required_modalities = ("rgb", "depth", "cam_pose")
404
405    def build_entries(self, df) -> List[Dict[str, Any]]:
406        rows = super().build_entries(df)
407        # Augment each paired entry with depth_a / depth_b.
408        for e in rows:
409            scene = e["scene_id"]
410            phase = e["phase"]
411            stem_a = e["metadata"]["frame"]
412            stem_b = e["metadata"]["frame_b"]
413            e["depth"] = _modality_path(scene, phase, "depth", stem_a)
414            e["depth_b"] = _modality_path(scene, phase, "depth", stem_b)
415        return rows
416
417
418class _ObjectTrackingSpec(_TaskSpec):
419    """``object_tracking``: RGB plus temporally consistent instance masks.
420
421    A mask value greater than zero is the persistent instance ID within one
422    ``(scene, phase)`` clip.  The production tracking runner groups these
423    frame entries into clips and uses the first mask as model initialisation.
424    No separate tracklet JSON is needed (and the pinned RPX release does not
425    publish one).
426    """
427
428    task_type_value = "object_tracking"
429    required_modalities = ("rgb", "masks")
430
431    def build_entries(self, df) -> List[Dict[str, Any]]:
432        rows: List[Dict[str, Any]] = []
433        for _, row in df.iterrows():
434            e = _base_entry(row)
435            scene = e["scene_id"]
436            phase = e["phase"]
437            stem = _frame_stem(row)
438            st = str(row["scene_type"])
439            e["rgb"] = _modality_path(scene, phase, "rgb", stem, scene_type=st)
440            e["mask"] = _modality_path(scene, phase, "masks", stem, scene_type=st)
441            rows.append(e)
442        return rows
443
444
445class _EgoSegmentationSpec(_SegmentationSpec):
446    """``ego_segmentation``: same shape as ``segmentation`` (rgb + mask),
447    scoped to ego (egocentric/GoPro) scenes only. Reuses
448    TaskType.OBJECT_SEGMENTATION — same evaluation task, just a distinct
449    manifest key so a downloader can pull ego's slice without also
450    pulling mos's rig-camera segmentation data (and vice versa)."""
451
452    scene_types = ("ego",)
453
454
455class _EgoObjectTrackingSpec(_ObjectTrackingSpec):
456    """``ego_object_tracking``: same shape as ``object_tracking``, scoped
457    to ego scenes only. See _EgoSegmentationSpec for why this is a
458    separate manifest key rather than a new TaskType."""
459
460    scene_types = ("ego",)
461
462
463class _VQASpec(_TaskSpec):
464    """``vqa``: not yet wired — needs spatial_qa / questionnaire labels.
465
466    Returns no entries until the team's VQA generation pipeline lands
467    (tracked separately in `paper-submission/VQA_DESIGN.md`).
468    Logging an explicit warning so this is visible in the upload run.
469    """
470
471    task_type_value = "visual_grounding"
472    required_modalities = ("rgb",)
473
474    def build_entries(self, df) -> List[Dict[str, Any]]:
475        log.warning(
476            "vqa: spec exists but no entries produced — needs the team's "
477            "VQA label generation pipeline (see VQA_DESIGN.md). Skipping."
478        )
479        return []
480
481
482class _EgoVQASpec(_VQASpec):
483    """``ego_vqa``: same reserved-slot shape as ``vqa``, scoped to ego
484    scenes only. See _EgoSegmentationSpec for why this is a separate
485    manifest key rather than a new TaskType."""
486
487    scene_types = ("ego",)
488
489    def build_entries(self, df) -> List[Dict[str, Any]]:
490        log.warning(
491            "ego_vqa: spec exists but no entries produced — needs the team's "
492            "VQA label generation pipeline (see VQA_DESIGN.md). Skipping."
493        )
494        return []
495
496
497# Recipe-key → spec instance.
498_TASK_SPECS: Dict[str, _TaskSpec] = {
499    "monocular_depth": _SingleFrameDepthSpec(),
500    "video_depth": _VideoDepthSpec(),
501    "segmentation": _SegmentationSpec(),
502    "rgbd_segmentation": _RGBDSegmentationSpec(),
503    "stereo_depth": _StereoDepthSpec(),
504    "relative_pose": _RelativePoseSpec(),
505    "rgbd_relative_pose": _RGBDRelativePoseSpec(),
506    "object_tracking": _ObjectTrackingSpec(),
507    "ego_segmentation": _EgoSegmentationSpec(),
508    "ego_object_tracking": _EgoObjectTrackingSpec(),
509    "vqa": _VQASpec(),
510    "ego_vqa": _EgoVQASpec(),
511}
512
513
514# ---------------------------------------------------------------------- #
515# Helpers
516# ---------------------------------------------------------------------- #
517
518_SPLITS: Tuple[str, ...] = ("easy", "medium", "hard")
519
520
521def _frame_stem(row) -> str:
522    """Return the frame's stem (no extension)."""
523    return Path(str(row["frame_filename"])).stem
524
525
526def _base_entry(row) -> Dict[str, Any]:
527    """Common per-row fields: id, scene_id, phase, frame_idx, difficulty,
528    metadata. Task-specific specs add modality paths on top."""
529    scene = str(row["scene_id"])
530    phase = int(row["phase"])
531    stem = _frame_stem(row)
532    return {
533        "id": f"{scene}__{phase}__{stem}",
534        "scene_id": scene,
535        "phase": phase,
536        "frame_idx": int(row["frame_idx"]),
537        "difficulty": str(row["split"]),
538        "metadata": {
539            "scene_id": scene,
540            "phase_idx": phase,
541            "frame": stem,
542        },
543    }
544
545
546def _single_frame_entry(row, modalities: Iterable[str]) -> Dict[str, Any]:
547    """Build a single-frame entry where each modality's *entry key* is the
548    same as its *modality name* (only correct for monocular_depth where
549    the loader's keys are literally ``rgb`` + ``depth``)."""
550    e = _base_entry(row)
551    scene = e["scene_id"]
552    phase = e["phase"]
553    stem = e["metadata"]["frame"]
554    for m in modalities:
555        e[m] = _modality_path(scene, phase, m, stem)
556    return e
557
558
559def _parquet_path(staging_root: Path) -> Path:
560    cand = staging_root / "manifest" / "frames_v1.parquet"
561    if cand.is_file():
562        return cand
563    matches = list((staging_root / "manifest").glob("frames_*.parquet"))
564    if not matches:
565        from ..exceptions import DatasetError
566
567        raise DatasetError(
568            f"no frames Parquet under {staging_root}/manifest/",
569            hint="Run `dataset_hub.cli manifest` first to build the parquet from your captures.",
570        )
571    return matches[-1]
572
573
574def _scene_wise_splits(df):
575    """Mode-aggregate per-row split into per-scene tier."""
576    return df.groupby("scene_id")["split"].agg(
577        lambda s: s.dropna().value_counts().idxmax() if s.notna().any() else None
578    )
579
580
581def _filter_required(df, required_cols: Iterable[str]):
582    """Return rows that have every required modality column set to True."""
583    sub = df
584    for col in required_cols:
585        if col not in sub.columns:
586            return None
587        sub = sub[sub[col].fillna(False).astype(bool)]
588    return sub
589
590
591# ---------------------------------------------------------------------- #
592# Public entry point
593# ---------------------------------------------------------------------- #
594
595
596def write_split_manifests(
597    staging_root: str | Path,
598    *,
599    tasks: Optional[Iterable[str]] = None,
600    splits: Iterable[str] = _SPLITS,
601) -> Dict[Tuple[str, str], Path]:
602    """Write per-(task, split) manifest JSONs under
603    ``<staging_root>/manifests/<task>/<split>.json``.
604
605    Returns ``{(recipe_key, split): path}`` for every JSON written.
606    Empty (task, split) combinations (no rows after filtering, or specs
607    that intentionally emit nothing — e.g. VQA) are skipped silently.
608    """
609    try:
610        import pandas as pd  # noqa: PLC0415
611    except ImportError as e:
612        raise ImportError(
613            "pandas is required for split_manifests; install via `pip install pandas`."
614        ) from e
615
616    staging_root = Path(staging_root)
617    parquet = _parquet_path(staging_root)
618    df = pd.read_parquet(parquet)
619
620    # Authoritative scene-wise splits.
621    scene_split = _scene_wise_splits(df)
622    df = df.drop(columns=["split"]).merge(
623        scene_split.rename("split"),
624        left_on="scene_id",
625        right_index=True,
626    )
627
628    # Read per-modality extensions from current.json (v2+ datasets).
629    # Datasets without this field (v1) fall through to the hardcoded
630    # defaults in _MODALITY_LAYOUT.
631    extensions = _load_modality_extensions(staging_root)
632    layout_override = _layout_with_overrides(extensions) if extensions else None
633    if extensions:
634        log.info(
635            "split_manifests: using modality_extensions from current.json: %s",
636            ", ".join(f"{k}={v}" for k, v in sorted(extensions.items())),
637        )
638
639    chosen = list(tasks) if tasks is not None else list(_TASK_SPECS.keys())
640    manifests_dir = staging_root / "manifests"
641    written: Dict[Tuple[str, str], Path] = {}
642
643    global _active_layout
644    _active_layout = layout_override
645    try:
646        written = _write_loop(
647            chosen=chosen,
648            df=df,
649            splits=splits,
650            manifests_dir=manifests_dir,
651        )
652    finally:
653        _active_layout = None
654
655    return written
656
657
658def _write_loop(
659    *,
660    chosen: List[str],
661    df,
662    splits: Iterable[str],
663    manifests_dir: Path,
664) -> Dict[Tuple[str, str], Path]:
665    """The per-task per-split write loop, factored out so the layout
666    override at module level wraps it in try/finally cleanly."""
667    written: Dict[Tuple[str, str], Path] = {}
668    for recipe_key in chosen:
669        if recipe_key not in _TASK_SPECS:
670            log.warning("unknown task %r — skipping", recipe_key)
671            continue
672        spec = _TASK_SPECS[recipe_key]
673
674        # Scope to this spec's scene family FIRST — required_modalities
675        # alone (rgb+masks) is satisfied by both mos and ego rows, so
676        # without this an ego frame would qualify for "segmentation" too,
677        # not just "ego_segmentation".
678        sub_all = df[df["scene_type"].isin(spec.scene_types)]
679
680        # Filter to rows that have every required modality.
681        required_cols = [_has_col(m) for m in spec.required_modalities]
682        sub_all = _filter_required(sub_all, required_cols)
683        if sub_all is None or sub_all.empty:
684            log.info(
685                "no parquet rows satisfy %s's required modalities (%s); skipping all splits.",
686                recipe_key,
687                list(spec.required_modalities),
688            )
689            continue
690
691        for split in splits:
692            sub = sub_all[sub_all["split"].astype(str) == split]
693            if sub.empty:
694                log.info("no samples for task=%s split=%s; skipping", recipe_key, split)
695                continue
696
697            samples = spec.build_entries(sub)
698            if not samples:
699                # Spec opted out (e.g., VQA needs labels not present yet).
700                continue
701
702            payload = {
703                "task": spec.task_type_value,  # ← TaskType.value, not recipe_key
704                "split": split,
705                "root": None,  # filled at download time
706                "samples": samples,
707            }
708            out = manifests_dir / recipe_key / f"{split}.json"
709            out.parent.mkdir(parents=True, exist_ok=True)
710            out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
711            log.info("wrote %s (%d samples, task_type=%s)", out, len(samples), spec.task_type_value)
712            written[(recipe_key, split)] = out
713
714    return written
715
716
717__all__ = ["write_split_manifests"]