rpx_benchmark.dataset_hub.packer

  1"""Pack an on-disk capture tree into per-modality tar shards for HF upload.
  2
  3We deliberately tar each ``(scene, phase, modality)`` triple instead of
  4shipping loose files. Three reasons:
  5
  61. **Upload speed.** The full RPX dataset is ~890 GB across an estimated
  7   750k loose frames. HF (and git LFS underneath) handles thousands of
  8   ~400 MB tars far faster than hundreds of thousands of loose files.
  9
 102. **Selective download granularity.** ``snapshot_download`` filters at the
 11   *file* level. With per-modality tars, a single ``allow_pattern`` like
 12   ``scenes/*/*/rgb.tar`` pulls exactly the bytes a task needs and
 13   nothing more — the cache deduplicates across tasks automatically.
 14
 153. **Atomic re-issue of label updates.** Bumping a label version is a
 16   one-line edit (re-run the packer for that modality with
 17   ``--label-version v2``) plus an upload of the new tars; older versions
 18   stay in the repo for reproducibility.
 19
 20Output layout (under the staging directory)::
 21
 22    staging/
 23    ├── scenes/                                # multi-object captures
 24    │   └── <scene_id>/<phase>/
 25    │       ├── rgb.tar                   # raw (sealed once)
 26    │       ├── depth.tar                 # raw
 27    │       ├── fisheye.tar               # raw
 28    │       └── labels/
 29    │           ├── cam_pose/v1.tar       # versioned: SLAM may be re-run
 30    │           ├── masks/v1.tar
 31    │           ├── masks_aux/v1.tar
 32    │           └── sam2_meta/v1.tar
 33    ├── objects/                               # single-object captures (one phase)
 34    │   └── <object_id>/0/
 35    │       (same modality structure)
 36    └── ego/                                    # egocentric captures (one phase)
 37        └── <scene_id>/0/
 38            (same modality structure — rgb + masks/masks_aux only today;
 39            scene_id matches the sibling mos/ scene so splits/difficulty
 40            join without any extra lookup table)
 41
 42The packer is *content-stable*: tar member order is sorted, mtime is
 43zeroed, ownership is normalised — so re-running on the same source
 44produces byte-identical tars (and therefore identical SHA-256 hashes).
 45This matters for the manifest's integrity field and for HF's
 46content-addressed cache to deduplicate across re-runs.
 47"""
 48
 49from __future__ import annotations
 50
 51import hashlib
 52import json
 53import os
 54import re
 55import tarfile
 56from dataclasses import dataclass, field
 57from pathlib import Path
 58from typing import Dict, List, Optional
 59
 60from ..exceptions import ConfigError, DatasetError
 61from ..logging_utils import get_logger
 62from .recipes import (
 63    CAM_POSE,
 64    DEPTH,
 65    FISHEYE,
 66    MASKS,
 67    MASKS_AUX,
 68    RGB,
 69    SceneType,
 70)
 71from .scanner import SRC_SUBDIR_BY_TYPE, ScanResult
 72
 73log = get_logger(__name__)
 74
 75
 76# Modalities considered "raw" (immutable, no version embedded in path).
 77# rgb/depth/fisheye are sealed once captured — never re-released.
 78# cam_pose is *versioned* (treated as a label) because the SLAM pipeline
 79# may be re-run; that version bump should not invalidate the rest of the
 80# user's cache.
 81RAW_MODALITIES = frozenset({RGB, DEPTH, FISHEYE})
 82
 83# Mapping from scanner modality keys → packer output modality + label flag.
 84# Scanner keys come straight from the on-disk dir names, with sam2/ split
 85# into sam2/<sub>. The packer collapses these into a small, named set.
 86_MODALITY_PLAN: Dict[str, tuple[str, bool]] = {
 87    # scanner key                 (output modality name, is_label)
 88    "rgb": (RGB, False),
 89    "depth": (DEPTH, False),
 90    "fisheye": (FISHEYE, False),
 91    "cam_pose": (CAM_POSE, True),  # versioned label
 92    "sam2/masks": (MASKS, True),
 93    "sam2/_meta": ("sam2_meta", True),
 94    # The remaining sam2/* aux directories all collapse into masks_aux.
 95    "sam2/bbox_overlay": (MASKS_AUX, True),
 96    "sam2/contour_gt_masks": (MASKS_AUX, True),
 97    "sam2/dino_output": (MASKS_AUX, True),
 98    "sam2/masks_contour_with_hidden": (MASKS_AUX, True),
 99    "sam2/palette": (MASKS_AUX, True),
100    "sam2/rgb_and_mask": (MASKS_AUX, True),
101}
102
103
104@dataclass(frozen=True)
105class PackPlan:
106    """Knobs for the packer."""
107
108    src_root: Path
109    staging_root: Path
110    label_version: str = "v1"
111    overwrite: bool = False
112    compute_hash: bool = True
113
114
115@dataclass(frozen=True)
116class PackedShard:
117    """One tar shard the packer produced."""
118
119    repo_path: str  # path relative to staging_root, forward-slashes
120    scene_id: str
121    scene_type: SceneType  # disambiguates scene_id collisions across families
122                            # (ego scenes deliberately reuse their mos/
123                            # sibling's scene_id — see manifest.py's
124                            # _index_shards_by_phase / _frame_filenames_for)
125    phase: int
126    modality: str
127    is_label: bool
128    file_count: int
129    total_bytes: int
130    sha256: Optional[str] = None
131
132
133@dataclass(frozen=True)
134class PackResult:
135    """Aggregate result of one packer run."""
136
137    shards: List[PackedShard]
138    skipped_keys: List[str] = field(default_factory=list)
139
140    @property
141    def total_bytes(self) -> int:
142        return sum(s.total_bytes for s in self.shards)
143
144    @property
145    def total_files(self) -> int:
146        return sum(s.file_count for s in self.shards)
147
148
149# --------------------------------------------------------------------- #
150# Helpers
151# --------------------------------------------------------------------- #
152
153
154#: Repo-path root per scene type. Public (no leading underscore) — other
155#: modules (cli.py's manifest-rebuild-from-existing-tars path) need this
156#: same mapping and should import it rather than keep their own copy.
157SCENE_ROOT_BY_TYPE: Dict[SceneType, str] = {
158    SceneType.MULTI_OBJECT: "scenes",
159    SceneType.SINGLE_OBJECT: "objects",
160    # ego nests under the SAME scenes/<scene_id>/ dir as its mos sibling,
161    # as a phase-like "ego" segment alongside 0/1/2 -- not its own
162    # top-level root. See phase_segment() below.
163    SceneType.EGO: "scenes",
164}
165
166
167def _scene_root_for(scene_type: SceneType) -> str:
168    return SCENE_ROOT_BY_TYPE[scene_type]
169
170
171def phase_segment(scene_type: SceneType, phase_index: int) -> str:
172    """The path segment standing in for "phase" in a repo path.
173
174    mos/sos use the numeric phase (0/1/2/...). ego has exactly one
175    continuous capture per scene and nests under its mos sibling's own
176    scene directory, so it uses the literal segment "ego" instead of a
177    number -- distinguishing ``scenes/<scene>/ego/`` from
178    ``scenes/<scene>/0/``, ``.../1/``, etc. Single source of truth: every
179    module that builds a repo path from (scene_type, phase) must call
180    this rather than stringifying phase_index directly, or ego paths
181    will end up as ``scenes/<scene>/0/...`` and collide with mos.
182    """
183    return "ego" if scene_type is SceneType.EGO else str(phase_index)
184
185
186def _shard_repo_path(
187    scene_type: SceneType,
188    scene_id: str,
189    phase: int,
190    modality: str,
191    is_label: bool,
192    label_version: str,
193) -> str:
194    base = f"{_scene_root_for(scene_type)}/{scene_id}/{phase_segment(scene_type, phase)}"
195    if is_label:
196        return f"{base}/labels/{modality}/{label_version}.tar"
197    return f"{base}/{modality}.tar"
198
199
200def _scene_src_root(src_root: Path, scene_type: SceneType, scene_id: str) -> Path:
201    """Where a scene actually lives on disk: ``<src>/{mos,sos,ego}/<scene_id>``."""
202    return src_root / SRC_SUBDIR_BY_TYPE[scene_type] / scene_id
203
204
205def _src_dir_for(
206    src_root: Path, scene_type: SceneType, scene_id: str, phase: int, scanner_key: str
207) -> Path:
208    """Translate a scanner modality key back to its on-disk source dir."""
209    return _scene_src_root(src_root, scene_type, scene_id) / str(phase) / scanner_key
210
211
212def _enumerate_files_for_modality(
213    src_root: Path,
214    scene_type: SceneType,
215    scene_id: str,
216    phase: int,
217    scanner_key: str,
218) -> List[Path]:
219    """All files under one modality dir, sorted, recursive.
220
221    For ``sam2/_meta`` the source is the ``sam2/`` dir itself, but only
222    its *immediate* loose files (not subdirs) are taken — the subdirs
223    are owned by the per-aux-dir scanner keys.
224    """
225    if scanner_key == "sam2/_meta":
226        sam2_root = _scene_src_root(src_root, scene_type, scene_id) / str(phase) / "sam2"
227        if not sam2_root.is_dir():
228            return []
229        return sorted(p for p in sam2_root.iterdir() if p.is_file())
230
231    src_dir = _src_dir_for(src_root, scene_type, scene_id, phase, scanner_key)
232    if not src_dir.is_dir():
233        return []
234    return sorted(p for p in src_dir.rglob("*") if p.is_file())
235
236
237def _write_tar(
238    out_path: Path,
239    files: List[Path],
240    rel_to: Path,
241    compute_hash: bool,
242) -> tuple[int, int, Optional[str]]:
243    """Write ``files`` to a tar at ``out_path``. Returns
244    ``(file_count, total_bytes, sha256)``. Tar entries are deterministic.
245    """
246    out_path.parent.mkdir(parents=True, exist_ok=True)
247    total = 0
248    h = hashlib.sha256() if compute_hash else None
249
250    # Write to a temp path then rename — never expose a half-written tar.
251    tmp = out_path.with_suffix(out_path.suffix + ".part")
252    # dereference=True: gettarinfo() must use os.stat() (follow symlinks),
253    # not os.lstat(). Without it, any symlinked source file -- e.g. ego
254    # scenes arranged via ego_layout.py's per-file symlinks -- would be
255    # packed as a broken symlink tar entry pointing at a local machine
256    # path (near-zero bytes, unreadable by anyone downloading the shard)
257    # instead of the real file content. See lossless_convert.py's
258    # _link_one docstring for how this was originally discovered.
259    with tarfile.open(tmp, mode="w", format=tarfile.USTAR_FORMAT, dereference=True) as tf:
260        for f in files:
261            arcname = str(f.relative_to(rel_to)).replace(os.sep, "/")
262            ti = tf.gettarinfo(name=str(f), arcname=arcname)
263            # Normalise so re-runs produce byte-identical tars.
264            ti.mtime = 0
265            ti.uid = 0
266            ti.gid = 0
267            ti.uname = ""
268            ti.gname = ""
269            with f.open("rb") as fh:
270                tf.addfile(ti, fh)
271            total += f.stat().st_size
272
273    if h is not None:
274        with tmp.open("rb") as fh:
275            for chunk in iter(lambda: fh.read(1 << 20), b""):
276                h.update(chunk)
277
278    os.replace(tmp, out_path)
279    return len(files), total, (h.hexdigest() if h else None)
280
281
282# --------------------------------------------------------------------- #
283# Public API
284# --------------------------------------------------------------------- #
285
286
287def pack_capture_tree(plan: PackPlan, scan: ScanResult) -> PackResult:
288    """Pack each ``(scene, phase, modality)`` triple in ``scan`` into one tar.
289
290    Reusing :class:`ScanResult` from the scanner avoids re-walking the
291    source tree — the scan already has the file counts, byte totals, and
292    structure we need. The packer just translates each (scene, phase,
293    modality) entry into a tar at the right repo path.
294    """
295    if plan.src_root != scan.root:
296        raise ConfigError(
297            f"plan.src_root ({plan.src_root}) does not match the root the "
298            f"scan was taken at ({scan.root}).",
299            hint="Pass the same root to scan_capture_root() and PackPlan.",
300        )
301    plan.staging_root.mkdir(parents=True, exist_ok=True)
302
303    shards: List[PackedShard] = []
304    skipped: List[str] = []
305
306    for scene in scan.scenes:
307        for phase in scene.phases:
308            # Group source modalities by output (modality, is_label) so
309            # masks_aux/* all flow into the same masks_aux.tar.
310            grouped: Dict[tuple[str, bool], list[str]] = {}
311            for scanner_key in phase.modalities:
312                target = _MODALITY_PLAN.get(scanner_key)
313                if target is None:
314                    skipped.append(f"{scene.scene_id}/{phase.phase_index}/{scanner_key}")
315                    continue
316                grouped.setdefault(target, []).append(scanner_key)
317
318            for (out_modality, is_label), keys in grouped.items():
319                files: List[Path] = []
320                for k in keys:
321                    files.extend(
322                        _enumerate_files_for_modality(
323                            plan.src_root,
324                            scene.scene_type,
325                            scene.scene_id,
326                            phase.phase_index,
327                            k,
328                        )
329                    )
330                if not files:
331                    continue
332
333                repo_path = _shard_repo_path(
334                    scene.scene_type,
335                    scene.scene_id,
336                    phase.phase_index,
337                    out_modality,
338                    is_label,
339                    plan.label_version,
340                )
341                out_path = plan.staging_root / repo_path
342                if out_path.exists() and not plan.overwrite:
343                    raise DatasetError(
344                        f"refusing to overwrite existing shard: {out_path}",
345                        hint="Pass overwrite=True or remove the staging dir.",
346                    )
347
348                fc, bs, sha = _write_tar(
349                    out_path,
350                    files,
351                    rel_to=(
352                        _scene_src_root(plan.src_root, scene.scene_type, scene.scene_id)
353                        / str(phase.phase_index)
354                    ),
355                    compute_hash=plan.compute_hash,
356                )
357                shards.append(
358                    PackedShard(
359                        repo_path=repo_path,
360                        scene_id=scene.scene_id,
361                        scene_type=scene.scene_type,
362                        phase=phase.phase_index,
363                        modality=out_modality,
364                        is_label=is_label,
365                        file_count=fc,
366                        total_bytes=bs,
367                        sha256=sha,
368                    )
369                )
370                log.info("packed %s (%d files, %d bytes)", repo_path, fc, bs)
371
372    return PackResult(shards=shards, skipped_keys=skipped)
373
374
375# --------------------------------------------------------------------- #
376# Per-object shared artefacts (objects_meta/) — questionnaire dedup
377# --------------------------------------------------------------------- #
378
379# FewSOL questionnaire is "<n>. <question>" then a comma-separated answer
380# block on the following non-blank line. This regex captures one Q-A pair.
381_QUESTION_RE = re.compile(r"^\s*\d+\.\s*(?P<q>.+\?)\s*$", re.MULTILINE)
382
383
384@dataclass(frozen=True)
385class SharedArtefact:
386    """One file written under ``objects_meta/<object_id>/``."""
387
388    object_id: str
389    repo_path: str
390    total_bytes: int
391
392
393def _parse_questionnaire(text: str) -> Dict[str, list[str]]:
394    """Parse the FewSOL-style questionnaire into a structured dict.
395
396    The text shape is::
397
398        # comments are ignored
399        1. What is the name of the object in these images?
400        tape and tape holder, tape, support and adhesive tape
401
402        2. What is the category of the object in these images?
403        ...
404
405    Returns ``{question_text: [answer1, answer2, ...]}`` with answers
406    split on commas. Questions are kept verbatim (warts and all) so the
407    schema survives changes to the answer-template wording.
408    """
409    out: Dict[str, list[str]] = {}
410    lines = [ln for ln in text.splitlines() if not ln.lstrip().startswith("#")]
411    cleaned = "\n".join(lines)
412    matches = list(_QUESTION_RE.finditer(cleaned))
413    for i, m in enumerate(matches):
414        q = m.group("q").strip()
415        start = m.end()
416        end = matches[i + 1].start() if i + 1 < len(matches) else len(cleaned)
417        body = cleaned[start:end].strip()
418        # Take the first non-blank line as the answer block.
419        first_line = next(
420            (ln.strip() for ln in body.splitlines() if ln.strip()),
421            "",
422        )
423        answers = [a.strip() for a in first_line.split(",") if a.strip()]
424        out[q] = answers
425    return out
426
427
428def pack_objects_meta(plan: PackPlan, scan: ScanResult) -> List[SharedArtefact]:
429    """Build the ``objects_meta/`` layer.
430
431    For every SOS scene with a ``questionnaire.txt`` at its scene root,
432    parse the file and write a normalised
433    ``objects_meta/<object_id>/questionnaire.json`` under the staging
434    dir. ``object_id`` is exactly the SOS scene directory name (the
435    canonical key MOS scenes use in ``mask_to_object.json``).
436
437    Also write ``objects_meta/_index.json``: a small lookup table
438    listing every known object_id. The downloader uses this to verify
439    that mask_to_object references resolve.
440    """
441    out: List[SharedArtefact] = []
442    objects_meta = plan.staging_root / "objects_meta"
443    objects_meta.mkdir(parents=True, exist_ok=True)
444
445    object_ids: list[str] = []
446    for scene in scan.scenes:
447        if scene.scene_type is not SceneType.SINGLE_OBJECT:
448            continue
449        sos_root = plan.src_root / "sos" / scene.scene_id
450        questionnaire_src = sos_root / "questionnaire.txt"
451        if not questionnaire_src.is_file():
452            log.warning("SOS scene %s has no questionnaire.txt; skipping", scene.scene_id)
453            continue
454
455        parsed = _parse_questionnaire(
456            questionnaire_src.read_text(encoding="utf-8"),
457        )
458        payload = {
459            "object_id": scene.scene_id,
460            "source": "sos/" + scene.scene_id + "/questionnaire.txt",
461            "questions": parsed,
462        }
463        repo_path = f"objects_meta/{scene.scene_id}/questionnaire.json"
464        out_path = plan.staging_root / repo_path
465        out_path.parent.mkdir(parents=True, exist_ok=True)
466        if out_path.exists() and not plan.overwrite:
467            raise DatasetError(
468                f"refusing to overwrite existing object meta: {out_path}",
469                hint="Pass overwrite=True or remove the staging dir.",
470            )
471        rendered = json.dumps(payload, indent=2, sort_keys=True)
472        out_path.write_text(rendered, encoding="utf-8")
473        out.append(
474            SharedArtefact(
475                object_id=scene.scene_id,
476                repo_path=repo_path,
477                total_bytes=out_path.stat().st_size,
478            )
479        )
480        object_ids.append(scene.scene_id)
481        log.info("wrote %s (%d questions)", repo_path, len(parsed))
482
483    index_path = objects_meta / "_index.json"
484    index_path.write_text(
485        json.dumps({"object_ids": sorted(object_ids)}, indent=2, sort_keys=True),
486        encoding="utf-8",
487    )
488    return out