rpx_benchmark.dataset_hub.manifest

  1"""Build the per-frame Parquet manifest the dataset hub ships.
  2
  3The manifest is the *single source of truth* for what's in the dataset:
  4one row per frame, with boolean columns marking which modalities are
  5present, plus split assignment, scene type, and pointers to the tar
  6shards that contain the actual bytes. It is shipped at
  7``manifest/frames_<schema_version>.parquet`` on the HF repo and is
  8small enough (~30 MB for 750k frames) that the downloader pulls it
  9eagerly before deciding what else to fetch.
 10
 11Schema (column / dtype / meaning)::
 12
 13    scene_id        str          "scene1.library.fountain" or "object001.mug"
 14    scene_type      str          "multi_object" | "single_object"
 15    phase           int32        0/1/2 for multi, 0 for single
 16    frame_idx       int32        0-based frame index inside the phase
 17    frame_filename  str          the source filename (e.g. "00000.png")
 18    split           str | null   "easy" | "medium" | "hard" (null if no split)
 19    has_rgb         bool         True iff the frame appears in rgb.tar
 20    has_depth       bool
 21    has_fisheye     bool
 22    has_cam_pose    bool
 23    has_masks       bool
 24    has_masks_aux   bool
 25    has_sam2_meta   bool
 26    shard_rgb       str | null   repo-relative path to rgb.tar (or null)
 27    shard_depth     str | null
 28    shard_fisheye   str | null
 29    shard_cam_pose  str | null
 30    shard_masks     str | null
 31    shard_masks_aux str | null
 32    shard_sam2_meta str | null
 33
 34Plus a sibling ``manifest/current.json`` that maps logical label names
 35to the version currently considered "default" (e.g. ``masks → v1``).
 36The downloader reads this first to resolve recipe modalities to actual
 37shard paths.
 38"""
 39
 40from __future__ import annotations
 41
 42import json
 43from dataclasses import dataclass
 44from pathlib import Path
 45from typing import Any, Dict, List, Mapping, Optional
 46
 47from ..exceptions import DatasetError
 48from ..logging_utils import get_logger
 49from .packer import PackedShard, PackResult
 50from .recipes import (
 51    CAM_POSE,
 52    DEPTH,
 53    FISHEYE,
 54    MASKS,
 55    MASKS_AUX,
 56    RGB,
 57    SceneType,
 58)  # noqa: F401  CAM_POSE used in default label_versions
 59from .scanner import SRC_SUBDIR_BY_TYPE, ScanResult
 60
 61log = get_logger(__name__)
 62
 63
 64SCHEMA_VERSION = "v1"
 65
 66# Modality → ("has_X", "shard_X") column-name pair. Order is the canonical
 67# column order in the Parquet table.
 68_COLUMNS: Dict[str, tuple[str, str]] = {
 69    RGB: ("has_rgb", "shard_rgb"),
 70    DEPTH: ("has_depth", "shard_depth"),
 71    FISHEYE: ("has_fisheye", "shard_fisheye"),
 72    CAM_POSE: ("has_cam_pose", "shard_cam_pose"),
 73    MASKS: ("has_masks", "shard_masks"),
 74    MASKS_AUX: ("has_masks_aux", "shard_masks_aux"),
 75    "sam2_meta": ("has_sam2_meta", "shard_sam2_meta"),
 76}
 77
 78
 79@dataclass(frozen=True)
 80class ManifestPaths:
 81    """Where the manifest writer dropped its outputs."""
 82
 83    parquet_path: Path
 84    current_json_path: Path
 85
 86
 87def _arrow():
 88    """Lazy-import pyarrow with a friendly error message."""
 89    try:
 90        import pyarrow as pa  # noqa: WPS433
 91        import pyarrow.parquet as pq  # noqa: WPS433
 92    except ImportError as e:
 93        raise DatasetError(
 94            "pyarrow is required to write the dataset manifest.",
 95            hint="Install with: pip install 'rpx-benchmark[hub]'",
 96        ) from e
 97    return pa, pq
 98
 99
100def _index_shards_by_phase(
101    pack: PackResult,
102) -> Dict[tuple[SceneType, str, int], Dict[str, PackedShard]]:
103    """Group shards by (scene_type, scene_id, phase) for fast lookup while
104    walking frames.
105
106    scene_type is part of the key deliberately — ego scenes intentionally
107    share their exact scene_id string with their mos/ sibling (that's how
108    split/difficulty assignment joins them), so (scene_id, phase) alone is
109    NOT unique once ego scenes exist. Keying by scene_id alone here used to
110    silently let one scene type's shards clobber the other's for any
111    (scene_id, phase) pair they both used — same bug class as
112    _frame_filenames_for below.
113    """
114    out: Dict[tuple[SceneType, str, int], Dict[str, PackedShard]] = {}
115    for s in pack.shards:
116        key = (s.scene_type, s.scene_id, s.phase)
117        out.setdefault(key, {})[s.modality] = s
118    return out
119
120
121def _frame_filenames_for(
122    scan: ScanResult,
123    scene_type: SceneType,
124    scene_id: str,
125    phase_index: int,
126) -> List[str]:
127    """Pick one modality (rgb, then depth, then anything) and return its
128    sorted filenames as the canonical frame list for that phase.
129
130    scene_type disambiguates scenes that share a scene_id string across
131    families (ego scenes deliberately reuse their mos/ sibling's scene_id —
132    see _index_shards_by_phase) — without it, `next(... if s.scene_id ==
133    scene_id)` could resolve to the WRONG scene's phase entirely.
134    """
135    scene = next(
136        s for s in scan.scenes if s.scene_id == scene_id and s.scene_type is scene_type
137    )
138    phase = next(p for p in scene.phases if p.phase_index == phase_index)
139    src = scan.root / SRC_SUBDIR_BY_TYPE[scene_type] / scene_id / str(phase_index)
140    for prefer in ("rgb", "depth", "fisheye"):
141        if prefer in phase.modalities:
142            return sorted(p.name for p in (src / prefer).iterdir() if p.is_file())
143    return []
144
145
146# Logical-modality → on-disk sub-path relative to a phase root. These map
147# from the split-manifest's modality vocabulary onto where the actual
148# files live in the scanned tree. Used by :func:`_detect_modality_extensions`
149# to sniff per-modality file extensions for ``current.json``.
150_MODALITY_SUBPATH_FOR_DETECTION: Dict[str, str] = {
151    "rgb": "rgb",
152    "depth": "depth",
153    "fisheye": "fisheye",
154    "fisheye_left": "fisheye/left",
155    "fisheye_right": "fisheye/right",
156    "masks": "sam2/masks",
157    "cam_pose": "cam_pose",
158}
159
160
161def _detect_modality_extensions(scan: ScanResult) -> Dict[str, str]:
162    """Sniff one file per modality directory in the scanned tree and
163    return ``{modality_name: ".png" | ".webp" | ".npz" | ".npy"}``.
164
165    Records the dataset's *on-disk* extension per modality so the
166    split-manifest writer (and any other downstream consumer) can build
167    paths that resolve after extraction without hardcoded assumptions.
168
169    Backward-compatible: if a modality's directory is missing or empty
170    in the scanned tree, that modality is omitted from the returned
171    dict, and downstream code falls back to its built-in defaults.
172
173    Detection runs per scene_type first, then merges: ``current.json``'s
174    ``modality_extensions`` is a single flat ``{modality: ext}`` dict
175    shared by every scene family (mos/sos/ego), with no scene_type
176    dimension. If two families actually disagree on a modality's on-disk
177    extension at manifest-build time (e.g. mos ``rgb`` already
178    lossless-converted to ``.webp`` while a freshly-added ego family is
179    still raw ``.png``), silently picking whichever scene the scan
180    happens to hit first would corrupt the OTHER family's manifest paths
181    without any error. Raise instead of guessing.
182    """
183    per_type: Dict[SceneType, Dict[str, str]] = {}
184    for scene in scan.scenes:
185        if not scene.phases:
186            continue
187        by_modality = per_type.setdefault(scene.scene_type, {})
188        if len(by_modality) == len(_MODALITY_SUBPATH_FOR_DETECTION):
189            continue  # this scene_type is already fully covered
190        sub = SRC_SUBDIR_BY_TYPE[scene.scene_type]
191        phase_root = scan.root / sub / scene.scene_id / str(scene.phases[0].phase_index)
192        for modality, subpath in _MODALITY_SUBPATH_FOR_DETECTION.items():
193            if modality in by_modality:
194                continue
195            mod_dir = phase_root / subpath
196            if not mod_dir.is_dir():
197                continue
198            # Pick the first regular file under this modality's subdir;
199            # rglob handles the case where files live one level deeper
200            # (e.g. ``fisheye/`` has ``left/`` and ``right/`` subdirs but
201            # no files directly inside).
202            for entry in sorted(mod_dir.rglob("*")):
203                if entry.is_file():
204                    by_modality[modality] = entry.suffix
205                    break
206
207    found: Dict[str, str] = {}
208    conflicts: List[str] = []
209    for scene_type, by_modality in per_type.items():
210        for modality, ext in by_modality.items():
211            if modality in found and found[modality] != ext:
212                conflicts.append(
213                    f"{modality!r} is {found[modality]!r} elsewhere but "
214                    f"{ext!r} for scene_type={scene_type.value!r}"
215                )
216                continue
217            found[modality] = ext
218
219    if conflicts:
220        raise DatasetError(
221            "Conflicting on-disk file extensions for the same modality "
222            f"across scene types: {'; '.join(conflicts)}.",
223            hint=(
224                "modality_extensions in current.json has no scene_type "
225                "dimension -- every family must agree on the same on-disk "
226                "extension per modality before the manifest is rebuilt. "
227                "Run lossless_convert over whichever family is out of sync "
228                "(commonly: a newly-added family that hasn't been "
229                "converted yet) and retry."
230            ),
231        )
232    return found
233
234
235def build_frame_manifest(
236    scan: ScanResult,
237    pack: PackResult,
238    out_dir: Path,
239    splits: Optional[Mapping[str, str]] = None,
240    label_versions: Optional[Mapping[str, str]] = None,
241) -> ManifestPaths:
242    """Write ``manifest/frames_<v>.parquet`` and ``manifest/current.json``.
243
244    Parameters
245    ----------
246    scan : ScanResult
247        Output of :func:`scan_capture_root`.
248    pack : PackResult
249        Output of :func:`pack_capture_tree`.
250    out_dir : Path
251        Where to write ``frames_<v>.parquet`` and ``current.json``. Both
252        end up under ``out_dir / "manifest" / ...``.
253    splits : Mapping[str, str], optional
254        ``{scene_id: "easy" | "medium" | "hard"}`` for multi-object scenes.
255        Single-object scenes are always assigned ``None``.
256    label_versions : Mapping[str, str], optional
257        Logical-label → version map written verbatim to ``current.json``.
258        Defaults to ``{"masks": "v1", "masks_aux": "v1", "sam2_meta": "v1"}``.
259    """
260    pa, pq = _arrow()
261    splits = dict(splits or {})
262    label_versions = dict(
263        label_versions
264        or {
265            MASKS: "v1",
266            MASKS_AUX: "v1",
267            "sam2_meta": "v1",
268            CAM_POSE: "v1",
269        }
270    )
271    out_dir = Path(out_dir)
272    (out_dir / "manifest").mkdir(parents=True, exist_ok=True)
273
274    shard_index = _index_shards_by_phase(pack)
275
276    # Build the table column-by-column for memory efficiency.
277    cols: Dict[str, list] = {
278        "scene_id": [],
279        "scene_type": [],
280        "phase": [],
281        "frame_idx": [],
282        "frame_filename": [],
283        "split": [],
284    }
285    for has_col, shard_col in _COLUMNS.values():
286        cols[has_col] = []
287        cols[shard_col] = []
288
289    for scene in scan.scenes:
290        for phase in scene.phases:
291            shards_here = shard_index.get(
292                (scene.scene_type, scene.scene_id, phase.phase_index), {}
293            )
294            filenames = _frame_filenames_for(
295                scan,
296                scene.scene_type,
297                scene.scene_id,
298                phase.phase_index,
299            )
300            # Ego scenes reuse their MOS sibling's split/difficulty tier —
301            # same scene_id (e.g. "scene20.su.checkerboard"), no separate
302            # ego difficulty scoring. SOS (single-object, no scene-level
303            # difficulty concept) still gets None.
304            split_label = (
305                splits.get(scene.scene_id)
306                if scene.scene_type in (SceneType.MULTI_OBJECT, SceneType.EGO)
307                else None
308            )
309            for idx, fname in enumerate(filenames):
310                cols["scene_id"].append(scene.scene_id)
311                cols["scene_type"].append(scene.scene_type.value)
312                cols["phase"].append(phase.phase_index)
313                cols["frame_idx"].append(idx)
314                cols["frame_filename"].append(fname)
315                cols["split"].append(split_label)
316                for modality, (has_col, shard_col) in _COLUMNS.items():
317                    shard = shards_here.get(modality)
318                    cols[has_col].append(shard is not None)
319                    cols[shard_col].append(shard.repo_path if shard else None)
320
321    schema = pa.schema(
322        [
323            ("scene_id", pa.string()),
324            ("scene_type", pa.string()),
325            ("phase", pa.int32()),
326            ("frame_idx", pa.int32()),
327            ("frame_filename", pa.string()),
328            ("split", pa.string()),
329            *[(c, pa.bool_()) for (c, _) in _COLUMNS.values()],
330            *[(c, pa.string()) for (_, c) in _COLUMNS.values()],
331        ]
332    )
333    table = pa.table({k: cols[k] for k in schema.names}, schema=schema)
334
335    parquet_path = out_dir / "manifest" / f"frames_{SCHEMA_VERSION}.parquet"
336    pq.write_table(table, parquet_path, compression="zstd")
337
338    current_path = out_dir / "manifest" / "current.json"
339    modality_extensions = _detect_modality_extensions(scan)
340
341    # Preserve any pre-existing current.json keys we do not own. The live
342    # IRVLUTD/RPX repo's current.json carries additional top-level blocks
343    # (``manifests``, ``sos``, ``mos``, ``metadata_versions``, …) that
344    # the teammate's separate manifest tooling writes. A naive overwrite
345    # would strip them and break downstream consumers; instead we read
346    # the existing payload (if any), update only the keys we are
347    # authoritative for, and write the merged result back.
348    #
349    # The keys we own (and will overwrite on every run):
350    #
351    #   * ``label_versions``       — from the ``--label-version`` flag
352    #   * ``schema_version``       — fixed at this module's SCHEMA_VERSION
353    #   * ``modality_extensions``  — sniffed from the scanned source tree
354    #
355    # Everything else is left untouched.
356    OWNED_KEYS = {"label_versions", "schema_version", "modality_extensions"}
357    if current_path.is_file():
358        try:
359            current_payload: Dict[str, Any] = json.loads(
360                current_path.read_text(encoding="utf-8"),
361            )
362            if not isinstance(current_payload, dict):
363                current_payload = {}
364        except json.JSONDecodeError:
365            current_payload = {}
366    else:
367        current_payload = {}
368    current_payload["label_versions"] = label_versions
369    current_payload["schema_version"] = SCHEMA_VERSION
370    if modality_extensions:
371        current_payload["modality_extensions"] = modality_extensions
372    # Don't carry stale modality_extensions if it has nothing to say
373    # AND no prior writer set it.
374    elif "modality_extensions" in current_payload and not current_payload["modality_extensions"]:
375        current_payload.pop("modality_extensions")
376    preserved = sorted(k for k in current_payload if k not in OWNED_KEYS)
377    if preserved:
378        log.info(
379            "current.json: preserved %d unmanaged top-level key(s): %s",
380            len(preserved),
381            ", ".join(preserved),
382        )
383    current_path.write_text(
384        json.dumps(current_payload, indent=2),
385        encoding="utf-8",
386    )
387
388    # Persist per-tar SHA-256 to manifest/checksums.json. The packer
389    # computes these as part of pack_capture_tree; the downstream
390    # download-side verification in hub._extract_snapshot_tars reads
391    # this file and aborts on any tar whose bytes don't match. Tar
392    # entries whose sha256 is None (rare — only when re-running
393    # manifest without a fresh pack) are still written with a null
394    # value so the absence is explicit, not silent.
395    checksums_path = out_dir / "manifest" / "checksums.json"
396    checksums = {s.repo_path: s.sha256 for s in pack.shards}
397    checksums_path.write_text(
398        json.dumps(
399            {"sha256": checksums, "algorithm": "sha256"},
400            indent=2,
401            sort_keys=True,
402        ),
403        encoding="utf-8",
404    )
405
406    log.info("wrote manifest: %d rows → %s", len(table), parquet_path)
407    n_with_sha = sum(1 for v in checksums.values() if v)
408    log.info(
409        "wrote checksums.json: %d/%d tar shards have SHA-256 → %s",
410        n_with_sha,
411        len(checksums),
412        checksums_path,
413    )
414
415    # Per-file SHA-256: walk every packed tar, stream each member through
416    # SHA-256, and persist the mapping at manifest/file_checksums.json.
417    # This is the foundation for ``verify_dataset(snapshot_root)`` to
418    # detect local-disk corruption between extract and load — the case
419    # the tar-level SHA-256 cannot guard against.
420    file_checksums = _compute_per_file_checksums(out_dir, pack)
421    file_checksums_path = out_dir / "manifest" / "file_checksums.json"
422    file_checksums_path.write_text(
423        json.dumps(
424            {"algorithm": "sha256", "sha256": file_checksums},
425            indent=2,
426            sort_keys=True,
427        ),
428        encoding="utf-8",
429    )
430    log.info(
431        "wrote file_checksums.json: %d extracted-file SHA-256 entries → %s",
432        len(file_checksums),
433        file_checksums_path,
434    )
435
436    if modality_extensions:
437        log.info(
438            "current.json modality_extensions: %s",
439            ", ".join(f"{k}={v}" for k, v in sorted(modality_extensions.items())),
440        )
441    return ManifestPaths(parquet_path=parquet_path, current_json_path=current_path)
442
443
444def _compute_per_file_checksums(
445    out_dir: Path,
446    pack: PackResult,
447) -> Dict[str, str]:
448    """For every member of every packed tar, compute SHA-256 and key it
449    by the *extracted* repo path the user will see after
450    ``_extract_snapshot_tars`` runs.
451
452    Returns ``{ "extracted/scenes/<scene>/<phase>/<member>": sha256_hex }``.
453
454    The per-tar SHA-256 (in ``checksums.json``) guards transport
455    corruption from operator → HF → user. This per-file map guards
456    everything *after* extraction: bit rot on the user's disk, OS-level
457    file corruption, accidental overwrites.
458    """
459    import hashlib
460    import tarfile
461
462    out: Dict[str, str] = {}
463    for shard in pack.shards:
464        tar_path = out_dir / shard.repo_path
465        if not tar_path.is_file():
466            continue
467        # The extracted layout mirrors the tar's nesting: a tar at
468        # ``scenes/scene_x/0/rgb.tar`` whose members are ``rgb/<frame>.webp``
469        # extracts to ``extracted/scenes/scene_x/0/rgb/<frame>.webp``.
470        # Strip the ``.tar`` filename to get the extracted prefix.
471        rel_parts = shard.repo_path.split("/")
472        if not rel_parts or not rel_parts[-1].endswith(".tar"):
473            continue
474        prefix_parts = rel_parts[:-1]
475        # labels/<modality>/<version>.tar extracts straight to
476        # extracted/.../<modality>'s file layout — drop the ``labels/`` and
477        # ``<version>`` segments so the result lives under the modality.
478        if "labels" in prefix_parts:
479            li = prefix_parts.index("labels")
480            prefix_parts = prefix_parts[:li]
481        extracted_prefix = "extracted/" + "/".join(prefix_parts) if prefix_parts else "extracted"
482        try:
483            with tarfile.open(tar_path, "r") as tf:
484                for member in tf.getmembers():
485                    if not member.isfile():
486                        continue
487                    f = tf.extractfile(member)
488                    if f is None:
489                        continue
490                    h = hashlib.sha256()
491                    for chunk in iter(lambda f=f: f.read(1 << 20), b""):
492                        h.update(chunk)
493                    extracted_path = f"{extracted_prefix}/{member.name}"
494                    out[extracted_path] = h.hexdigest()
495        except tarfile.TarError:
496            continue
497    return out
498
499
500def read_frame_manifest(parquet_path: Path):
501    """Read a manifest Parquet back as a pyarrow Table (test/debug aid)."""
502    _, pq = _arrow()
503    return pq.read_table(parquet_path)