rpx_benchmark.dataset_hub.cli

  1"""``python -m rpx_benchmark.dataset_hub.cli`` — sub-commands for the dataset hub.
  2
  3Available sub-commands::
  4
  5    mock      --out PATH [--multi N] [--single M] [--frames F] [--size S]
  6    scan      PATH                                                 [--json]
  7    pack      --src PATH --staging PATH [--label-version V] [--overwrite]
  8    manifest  --src PATH --pack-staging PATH --staging PATH [--splits FILE]
  9    upload    --staging PATH --repo-id ORG/REPO [--dry-run] [--public]
 10
 11The CLI uses argparse + sub-parsers so we can grow it without breaking
 12existing invocations.
 13"""
 14
 15from __future__ import annotations
 16
 17import argparse
 18import json
 19import os
 20import sys
 21from pathlib import Path
 22from typing import List, Optional
 23
 24from .croissant import CroissantPatch, stage_croissant
 25from .dataset_card import CardSpec, write_dataset_card
 26from .downloader import download_for_task
 27from .ego_layout import prepare_ego_layout
 28from .lossless_convert import ConvertSpec, convert_capture_tree
 29from .manifest import build_frame_manifest
 30from .mock import MockSpec, generate_mock, measure_tree
 31from .packer import PackPlan, pack_capture_tree, pack_objects_meta
 32from .recipes import DEFAULT_REPO_ID, SceneType
 33from .scanner import scan_capture_root
 34from .staging import load_scene_splits, stage_splits
 35from .uploader import UploadPlan, upload_staging
 36
 37# --------------------------------------------------------------------- #
 38# Formatting helpers
 39# --------------------------------------------------------------------- #
 40
 41
 42def _human_bytes(n: int) -> str:
 43    for unit in ("B", "KB", "MB", "GB", "TB"):
 44        if n < 1024 or unit == "TB":
 45            return f"{n:6.1f} {unit}" if unit != "B" else f"{n:6d} {unit}"
 46        n /= 1024
 47    return f"{n:.1f} TB"
 48
 49
 50# --------------------------------------------------------------------- #
 51# `mock`
 52# --------------------------------------------------------------------- #
 53
 54
 55def _cmd_mock(args: argparse.Namespace) -> int:
 56    spec = MockSpec(
 57        multi_object_scenes=args.multi,
 58        single_object_scenes=args.single,
 59        frames_per_phase=args.frames,
 60        image_size=args.size,
 61    )
 62    out = generate_mock(Path(args.out), spec)
 63    files, total = measure_tree(out)
 64    print(f"[mock] wrote {files} files ({_human_bytes(total)}) under {out}")
 65    return 0
 66
 67
 68# --------------------------------------------------------------------- #
 69# `scan`
 70# --------------------------------------------------------------------- #
 71
 72
 73def _cmd_scan(args: argparse.Namespace) -> int:
 74    res = scan_capture_root(Path(args.root))
 75
 76    if args.json:
 77        payload = {
 78            "root": str(res.root),
 79            "totals": {
 80                "scenes": len(res.scenes),
 81                "files": res.file_count,
 82                "bytes": res.total_bytes,
 83            },
 84            "by_type": {
 85                scene_type.value: {
 86                    "scenes": len(res.by_type(scene_type)),
 87                    "bytes": sum(s.total_bytes for s in res.by_type(scene_type)),
 88                }
 89                for scene_type in SceneType
 90            },
 91            "modalities": {
 92                name: {"files": inv.file_count, "bytes": inv.total_bytes}
 93                for name, inv in res.modality_totals().items()
 94            },
 95            "skipped": [str(p) for p in res.skipped],
 96        }
 97        json.dump(payload, sys.stdout, indent=2)
 98        sys.stdout.write("\n")
 99        return 0
100
101    print(f"\n📂  {res.root}")
102    print(
103        f"    {len(res.scenes)} scenes, {res.file_count:,} files, "
104        f"{_human_bytes(res.total_bytes)} total\n"
105    )
106
107    for scene_type in SceneType:
108        subset = res.by_type(scene_type)
109        if not subset:
110            continue
111        bytes_ = sum(s.total_bytes for s in subset)
112        files_ = sum(s.file_count for s in subset)
113        print(
114            f"  · {scene_type.value:14s}  {len(subset):4d} scenes  "
115            f"{files_:>9,} files  {_human_bytes(bytes_)}"
116        )
117
118    print("\n  per-modality totals (across all scenes & phases):")
119    for name, inv in res.modality_totals().items():
120        print(f"    {name:24s}  {inv.file_count:>9,} files  {_human_bytes(inv.total_bytes)}")
121
122    if res.skipped:
123        print(f"\n  ⚠ skipped {len(res.skipped)} non-matching entries:")
124        for p in res.skipped[:5]:
125            print(f"    - {p.name}")
126        if len(res.skipped) > 5:
127            print(f"    ... and {len(res.skipped) - 5} more")
128    print()
129    return 0
130
131
132def _cmd_prepare_ego(args: argparse.Namespace) -> int:
133    scenes = args.scenes.split(",") if args.scenes else None
134    report = prepare_ego_layout(
135        ego_captures_root=Path(args.ego_root),
136        mos_root=Path(args.mos_root),
137        dst_root=Path(args.dst),
138        scenes=scenes,
139        include_masks_verified=args.include_masks_verified,
140        mode="copy" if args.copy else "symlink",
141    )
142    ok = report.ok
143    skipped = report.skipped
144    print(f"\n  arranged {len(ok)} ego scene(s) -> {Path(args.dst) / 'ego'}")
145    for r in ok:
146        print(f"    ✓ {r.ego_scene_dir:12s} -> ego/{r.scene_id}/0")
147    if skipped:
148        print(f"\n  ⚠ skipped {len(skipped)}:")
149        for r in skipped:
150            print(f"    - {r.ego_scene_dir:12s} [{r.status}] {r.message}")
151    print(
152        "\n  NOTE: scene_id join to splits/scene_splits.json is an exact string "
153        "match — verify your mos_root's directory names actually match the "
154        "splits file's naming (padding/casing can silently drift; run `scan` "
155        "and cross-check before trusting difficulty tiers downstream)."
156    )
157    print()
158    return 0
159
160
161# --------------------------------------------------------------------- #
162# `pack`
163# --------------------------------------------------------------------- #
164
165
166def _cmd_pack(args: argparse.Namespace) -> int:
167    src = Path(args.src)
168    staging = Path(args.staging)
169    scan = scan_capture_root(src)
170    plan = PackPlan(
171        src_root=src,
172        staging_root=staging,
173        label_version=args.label_version,
174        overwrite=args.overwrite,
175    )
176    res = pack_capture_tree(plan, scan)
177    print(
178        f"[pack] wrote {len(res.shards)} tar shards "
179        f"({_human_bytes(res.total_bytes)} total) under {staging}"
180    )
181    if res.skipped_keys:
182        print(
183            f"[pack] skipped {len(res.skipped_keys)} unknown modality keys "
184            f"(first 3: {res.skipped_keys[:3]})"
185        )
186    artefacts = pack_objects_meta(plan, scan)
187    print(f"[pack] wrote {len(artefacts)} per-object questionnaires under {staging}/objects_meta/")
188    return 0
189
190
191# --------------------------------------------------------------------- #
192# `manifest`
193# --------------------------------------------------------------------- #
194
195
196def _cmd_manifest(args: argparse.Namespace) -> int:
197    src = Path(args.src)
198    staging = Path(args.staging)
199    scan = scan_capture_root(src)
200    # Re-pack-walk: rebuild PackResult by mirroring what the staging dir
201    # currently has on disk. We re-derive shard metadata from the tars.
202    pack = _rehydrate_pack_result(scan, staging)
203
204    splits = {}
205    if args.splits:
206        splits = json.loads(Path(args.splits).read_text(encoding="utf-8"))
207        # Accept both a flat {scene_id: tier} dict and the
208        # scene_splits.json shape {"splits": {"easy": [...], ...}}.
209        if "splits" in splits and isinstance(splits["splits"], dict):
210            splits = {sid: tier for tier, ids in splits["splits"].items() for sid in ids}
211
212    paths = build_frame_manifest(scan, pack, staging, splits=splits)
213    print(f"[manifest] wrote {paths.parquet_path}")
214    print(f"[manifest] wrote {paths.current_json_path}")
215
216    # Per-task per-split JSONs the toolkit's reader (`download_split`)
217    # expects at ``manifests/<task>/<split>.json``. Without these, a
218    # freshly published HF dataset 404s on every benchmark call.
219    from .split_manifests import write_split_manifests  # local import to avoid cycles
220
221    written = write_split_manifests(staging)
222    if not written:
223        # Loud failure: the published dataset would be unusable.
224        # Most common cause: ``--splits`` was not passed, so every row
225        # in the parquet has ``split=None`` and no per-task JSONs are
226        # produced. The hub README's canonical sequence requires
227        # ``--splits benchmark/data/splits/scene_splits.json``.
228        from ..exceptions import DatasetError
229
230        raise DatasetError(
231            "no per-task per-split manifests written — published HF tree "
232            "would be unusable for `download_split`. Most likely cause: "
233            "the --splits flag was not passed. Re-run with: "
234            "`--splits benchmark/data/splits/scene_splits.json` (or "
235            "another scene-tier mapping).",
236            hint="If you intentionally want to skip split-manifests for "
237            "a partial scene set, pass --splits with at least one "
238            "scene mapped to a tier.",
239        )
240    print(
241        f"[manifest] wrote {len(written)} per-task per-split manifests under {staging}/manifests/"
242    )
243    return 0
244
245
246def _rehydrate_pack_result(scan, staging):
247    """Best-effort reconstruction of PackResult from the staging dir.
248
249    Re-reads tar member counts and sizes from disk and **re-computes
250    each tar's SHA-256**. The packer writes deterministic tars, so the
251    recomputed hash is the same as the one the packer originally
252    produced. We need it persisted to ``checksums.json`` for
253    download-side verification — silently dropping it to ``None``
254    (as the previous implementation did) would have left the
255    fault-proof contract impossible to enforce at load time.
256
257    For the ~800 GB full-dataset pack on a modern NVMe disk this adds
258    a few minutes to the manifest step — a one-time, upload-machine
259    cost that buys per-tar tamper-evidence for every user from then on.
260    """
261    import hashlib
262    import tarfile
263
264    from .packer import (  # local import to avoid cycles
265        SCENE_ROOT_BY_TYPE,
266        PackedShard,
267        PackResult,
268        phase_segment,
269    )
270
271    def _file_sha256(path: Path) -> str:
272        h = hashlib.sha256()
273        with path.open("rb") as f:
274            for chunk in iter(lambda: f.read(1 << 20), b""):
275                h.update(chunk)
276        return h.hexdigest()
277
278    shards = []
279    for scene in scan.scenes:
280        for phase in scene.phases:
281            scene_root = SCENE_ROOT_BY_TYPE[scene.scene_type]
282            seg = phase_segment(scene.scene_type, phase.phase_index)
283            base = Path(staging) / scene_root / scene.scene_id / seg
284            for tar_path in sorted(base.rglob("*.tar")) if base.is_dir() else []:
285                try:
286                    with tarfile.open(tar_path, "r") as tf:
287                        members = tf.getmembers()
288                        fc = sum(1 for m in members if m.isfile())
289                except tarfile.TarError:
290                    fc = 0
291                rel = tar_path.relative_to(staging).as_posix()
292                # Modality name: rgb.tar -> "rgb"; labels/masks/v1.tar -> "masks"
293                if "/labels/" in rel:
294                    modality = rel.split("/labels/", 1)[1].split("/", 1)[0]
295                    is_label = True
296                else:
297                    modality = tar_path.stem
298                    is_label = False
299                shards.append(
300                    PackedShard(
301                        repo_path=rel,
302                        scene_id=scene.scene_id,
303                        scene_type=scene.scene_type,
304                        phase=phase.phase_index,
305                        modality=modality,
306                        is_label=is_label,
307                        file_count=fc,
308                        total_bytes=tar_path.stat().st_size,
309                        sha256=_file_sha256(tar_path),
310                    )
311                )
312    return PackResult(shards=shards)
313
314
315# --------------------------------------------------------------------- #
316# `upload`
317# --------------------------------------------------------------------- #
318
319
320def _cmd_upload(args: argparse.Namespace) -> int:
321    plan = UploadPlan(
322        staging_root=Path(args.staging),
323        repo_id=args.repo_id,
324        revision=args.revision,
325        private=not args.public,
326        dry_run=args.dry_run,
327        commit_message=args.message,
328    )
329    res = upload_staging(plan)
330    if res.dry_run if hasattr(res, "dry_run") else args.dry_run:
331        print(
332            f"[upload-dry-run] would push {res.files_planned} files "
333            f"({_human_bytes(res.bytes_planned)}) to {res.repo_id}"
334        )
335    else:
336        print(
337            f"[upload] pushed {res.files_planned} files "
338            f"({_human_bytes(res.bytes_planned)}) to {res.repo_id}"
339        )
340        if res.commit_url:
341            print(f"[upload] commit: {res.commit_url}")
342    return 0
343
344
345# --------------------------------------------------------------------- #
346# `stage-splits`
347# --------------------------------------------------------------------- #
348
349
350def _cmd_stage_splits(args: argparse.Namespace) -> int:
351    src = Path(args.splits_src) if args.splits_src else None
352    staged = stage_splits(
353        staging_root=Path(args.staging),
354        splits_src=src,
355        overwrite=args.overwrite,
356        require_all=not args.allow_missing,
357    )
358    print(f"[stage-splits] copied {len(staged)} files into {args.staging}/splits/")
359    for s in staged:
360        print(f"               {s.repo_path}  ({_human_bytes(s.bytes_)})")
361    return 0
362
363
364# --------------------------------------------------------------------- #
365# `dataset-card`
366# --------------------------------------------------------------------- #
367
368
369def _cmd_dataset_card(args: argparse.Namespace) -> int:
370    src = Path(args.src)
371    staging = Path(args.staging)
372    scan = scan_capture_root(src)
373
374    splits = None
375    if (staging / "splits" / "scene_splits.json").is_file():
376        try:
377            splits = load_scene_splits(staging)
378        except Exception as e:
379            print(f"[dataset-card] WARN: could not load scene splits: {e}", file=sys.stderr)
380
381    label_versions = None
382    cur_path = staging / "manifest" / "current.json"
383    if cur_path.is_file():
384        label_versions = json.loads(cur_path.read_text("utf-8")).get(
385            "label_versions",
386        )
387
388    spec = CardSpec(repo_id=args.repo_id)
389    out = write_dataset_card(
390        staging,
391        scan,
392        spec=spec,
393        splits=splits,
394        label_versions=label_versions,
395        overwrite=args.overwrite,
396    )
397    print(f"[dataset-card] wrote {out} ({out.stat().st_size:,} bytes)")
398    return 0
399
400
401# --------------------------------------------------------------------- #
402# `verify` — per-file SHA-256 integrity check of an extracted snapshot
403# --------------------------------------------------------------------- #
404
405
406def _cmd_verify(args: argparse.Namespace) -> int:
407    from ..exceptions import DatasetError
408    from ..hub import verify_dataset
409
410    try:
411        report = verify_dataset(Path(args.snapshot), raise_on_error=False)
412    except DatasetError as e:
413        print(f"[verify] {e}", file=sys.stderr)
414        return 2
415
416    if not report:
417        print(
418            f"[verify] no file_checksums.json under "
419            f"{args.snapshot}/manifest/ — legacy dataset, skipping per-file check.",
420            file=sys.stderr,
421        )
422        return 0
423
424    ok = sum(1 for v in report.values() if v == "ok")
425    bad = [(p, v) for p, v in report.items() if v != "ok"]
426    print(f"[verify] {ok} / {len(report)} files OK")
427    if not bad:
428        return 0
429    print(f"[verify] {len(bad)} failures:")
430    for p, status in sorted(bad)[:20]:
431        print(f"  {status:10s}  {p}")
432    if len(bad) > 20:
433        print(f"  ... and {len(bad) - 20} more")
434    return 1
435
436
437# --------------------------------------------------------------------- #
438# `lossless-convert`
439# --------------------------------------------------------------------- #
440
441
442def _cmd_lossless_convert(args: argparse.Namespace) -> int:
443    # verify=True is unconditional from the CLI. Per-frame round-trip
444    # verification is the load-bearing guarantee that the output tree
445    # is bit-faithful to the source; we do not expose any way to skip
446    # it from the published CLI surface. (The ConvertSpec field still
447    # exists for programmatic callers — tests, in particular — but
448    # operator-facing flags don't let it be set False.)
449    spec = ConvertSpec(
450        src_root=Path(args.src),
451        out_root=Path(args.out),
452        workers=args.workers,
453        dry_run=args.dry_run,
454        verify=True,
455        overwrite_out=args.overwrite_out,
456        webp_method=args.webp_method,
457        png_compress_level=args.png_compress_level,
458        skip_rgb_webp=args.skip_rgb_webp,
459        skip_png_recompress=args.skip_png_recompress,
460        skip_cam_pose=args.skip_cam_pose,
461    )
462    res = convert_capture_tree(spec)
463    if spec.dry_run:
464        print(
465            f"[lossless-convert] DRY-RUN: would convert {res.files_converted:,} files, "
466            f"link {res.files_linked:,} files (total {res.files_seen:,})"
467        )
468        for action, stats in sorted(res.by_action.items()):
469            print(f"  {action:18s} {stats.files:>7,} files")
470        return 0
471    # Per-action savings breakdown
472    for action, stats in sorted(res.by_action.items()):
473        if stats.files == 0:
474            continue
475        print(
476            f"[lossless-convert] {action:18s} "
477            f"{stats.files:>7,} files  "
478            f"{_human_bytes(stats.bytes_before)} → {_human_bytes(stats.bytes_after)}  "
479            f"saved {_human_bytes(stats.saved_bytes)} ({stats.saved_pct:.1f}%)"
480        )
481    print(
482        f"[lossless-convert] TOTAL: {res.files_seen:,} files  "
483        f"{_human_bytes(res.bytes_before)} → {_human_bytes(res.bytes_after)}  "
484        f"saved {_human_bytes(res.saved_bytes)} ({res.saved_pct:.1f}%)"
485    )
486    return 0
487
488
489# --------------------------------------------------------------------- #
490# `stage-croissant`
491# --------------------------------------------------------------------- #
492
493
494def _cmd_stage_croissant(args: argparse.Namespace) -> int:
495    src = Path(args.src) if args.src else None
496    patch = CroissantPatch(repo_id=args.repo_id, version=args.version, cite_as=args.cite_as)
497    out = stage_croissant(
498        staging_root=Path(args.staging), src=src, patch=patch, overwrite=args.overwrite
499    )
500    print(f"[stage-croissant] wrote {out} ({out.stat().st_size:,} bytes)")
501    return 0
502
503
504# --------------------------------------------------------------------- #
505# `download`
506# --------------------------------------------------------------------- #
507
508
509def _cmd_download(args: argparse.Namespace) -> int:
510    extras = tuple(args.modalities.split(",")) if args.modalities else ()
511    label_versions = {}
512    for kv in args.label_version or ():
513        k, _, v = kv.partition("=")
514        if not k or not v:
515            print(f"[download] ignoring malformed --label-version entry: {kv}", file=sys.stderr)
516            continue
517        label_versions[k] = v
518
519    res = download_for_task(
520        task=args.task,
521        split=args.split,
522        repo_id=args.repo_id,
523        revision=args.revision,
524        cache_dir=Path(args.cache_dir) if args.cache_dir else None,
525        extra_modalities=extras,
526        label_versions=label_versions,
527    )
528    print(f"[download] task={res.task} split={res.split} scene_type={res.scene_type.value}")
529    print(
530        f"[download] matched {len(res.matched_scenes)} scenes, {len(res.allow_patterns)} patterns"
531    )
532    print(f"[download] local_dir={res.local_dir}")
533    print(f"[download] {res.files_fetched} files on disk ({_human_bytes(res.bytes_fetched)})")
534    return 0
535
536
537# --------------------------------------------------------------------- #
538# Wiring
539# --------------------------------------------------------------------- #
540
541
542def build_parser() -> argparse.ArgumentParser:
543    parser = argparse.ArgumentParser(
544        prog="python -m rpx_benchmark.dataset_hub.cli",
545        description="RPX dataset-hub utilities (mock / scan).",
546    )
547    sub = parser.add_subparsers(dest="cmd", required=True)
548
549    p_mock = sub.add_parser("mock", help="Generate a synthetic capture tree.")
550    p_mock.add_argument("--out", required=True, help="Output root (will be created).")
551    p_mock.add_argument(
552        "--multi", type=int, default=3, help="Number of multi-object scenes (default: 3)."
553    )
554    p_mock.add_argument(
555        "--single", type=int, default=5, help="Number of single-object scenes (default: 5)."
556    )
557    p_mock.add_argument("--frames", type=int, default=4, help="Frames per phase (default: 4).")
558    p_mock.add_argument(
559        "--size", type=int, default=32, help="Image edge length in px (default: 32)."
560    )
561    p_mock.set_defaults(func=_cmd_mock)
562
563    p_scan = sub.add_parser("scan", help="Inventory an on-disk capture tree.")
564    p_scan.add_argument("root", help="Capture-tree root (test_dataset_aggregated/).")
565    p_scan.add_argument(
566        "--json", action="store_true", help="Emit machine-readable JSON instead of a table."
567    )
568    p_scan.set_defaults(func=_cmd_scan)
569
570    p_pe = sub.add_parser(
571        "prepare-ego",
572        help="Arrange raw ego (GoPro) captures into <DATA>/ego/<scene_id>/0/ for scan/pack.",
573    )
574    p_pe.add_argument("--ego-root", dest="ego_root", required=True,
575                      help="Root holding sceneNNN/ego/{rgb,sam2} captures.")
576    p_pe.add_argument("--mos-root", dest="mos_root", required=True,
577                      help="The mos/ dir to resolve each ego scene's sibling scene_id against.")
578    p_pe.add_argument("--dst", required=True,
579                      help="Capture-tree root to arrange into (writes <dst>/ego/...).")
580    p_pe.add_argument("--scenes", default=None,
581                      help="Comma-separated ego scene dir names to restrict to (default: all).")
582    p_pe.add_argument("--include-masks-verified", action="store_true",
583                      help="Include sam2/masks_verified/ (excluded by default — no MOS precedent).")
584    p_pe.add_argument("--copy", action="store_true",
585                      help="Copy files instead of symlinking (default: symlink, no duplication).")
586    p_pe.set_defaults(func=_cmd_prepare_ego)
587
588    p_pack = sub.add_parser("pack", help="Pack a capture tree into HF tar shards.")
589    p_pack.add_argument("--src", required=True, help="Capture-tree root.")
590    p_pack.add_argument("--staging", required=True, help="Staging output dir.")
591    p_pack.add_argument("--label-version", default="v1", help="Label version tag (default: v1).")
592    p_pack.add_argument(
593        "--overwrite", action="store_true", help="Overwrite existing shards in the staging dir."
594    )
595    p_pack.set_defaults(func=_cmd_pack)
596
597    p_man = sub.add_parser("manifest", help="Build the per-frame Parquet manifest.")
598    p_man.add_argument("--src", required=True, help="Capture-tree root (for frame inventory).")
599    p_man.add_argument(
600        "--staging", required=True, help="Packed staging dir (must contain *.tar shards)."
601    )
602    p_man.add_argument(
603        "--splits", default=None, help="Optional scene_splits.json path for split labels."
604    )
605    p_man.set_defaults(func=_cmd_manifest)
606
607    p_up = sub.add_parser("upload", help="Push a staging dir to a HF dataset repo.")
608    p_up.add_argument("--staging", required=True, help="Packed staging dir.")
609    p_up.add_argument(
610        "--repo-id",
611        default=DEFAULT_REPO_ID,
612        help=f"HF dataset repo id (default: {DEFAULT_REPO_ID}).",
613    )
614    p_up.add_argument("--revision", default="main", help="HF revision/branch (default: main).")
615    p_up.add_argument("--public", action="store_true", help="Create as public; default is private.")
616    p_up.add_argument(
617        "--dry-run", action="store_true", help="Walk the tree and report; no network call."
618    )
619    p_up.add_argument("--message", default="RPX dataset hub upload", help="Commit message.")
620    p_up.set_defaults(func=_cmd_upload)
621
622    p_dl = sub.add_parser("download", help="Pull just the tar shards a (task, split) needs.")
623    p_dl.add_argument(
624        "--task", required=True, help="Recipe key (e.g. segmentation, relative_pose)."
625    )
626    p_dl.add_argument("--split", default=None, help="easy | medium | hard (multi-object only).")
627    p_dl.add_argument(
628        "--repo-id",
629        default=DEFAULT_REPO_ID,
630        help=f"HF dataset repo id (default: {DEFAULT_REPO_ID}).",
631    )
632    p_dl.add_argument("--revision", default=None, help="HF revision/branch/tag (default: main).")
633    p_dl.add_argument(
634        "--cache-dir", default=None, help="HF cache root (default: ~/.cache/huggingface)."
635    )
636    p_dl.add_argument(
637        "--modalities", default="", help="Comma-separated extras beyond the recipe defaults."
638    )
639    p_dl.add_argument(
640        "--label-version",
641        action="append",
642        default=[],
643        metavar="KEY=VALUE",
644        help="Override a label modality version (repeatable).",
645    )
646    p_dl.set_defaults(func=_cmd_download)
647
648    p_st = sub.add_parser("stage-splits", help="Copy splits/* into a staging dir.")
649    p_st.add_argument("--staging", required=True, help="Staging dir (the upload target).")
650    p_st.add_argument(
651        "--splits-src",
652        default=None,
653        help="Source dir holding scene_splits.json + tier txts (default: benchmark/data/splits/).",
654    )
655    p_st.add_argument(
656        "--overwrite", action="store_true", help="Overwrite existing files in <staging>/splits/."
657    )
658    p_st.add_argument(
659        "--allow-missing",
660        action="store_true",
661        help="Don't raise if some expected splits files are absent in the source dir.",
662    )
663    p_st.set_defaults(func=_cmd_stage_splits)
664
665    p_dc = sub.add_parser("dataset-card", help="Generate <staging>/README.md (HF dataset card).")
666    p_dc.add_argument("--src", required=True, help="Capture-tree root (used for scan totals).")
667    p_dc.add_argument(
668        "--staging", required=True, help="Staging dir; the README lands at <staging>/README.md."
669    )
670    p_dc.add_argument(
671        "--repo-id",
672        default=DEFAULT_REPO_ID,
673        help=f"HF dataset repo id (default: {DEFAULT_REPO_ID}).",
674    )
675    p_dc.add_argument("--overwrite", action="store_true", help="Overwrite an existing README.md.")
676    p_dc.set_defaults(func=_cmd_dataset_card)
677
678    p_vfy = sub.add_parser(
679        "verify",
680        help=(
681            "Per-file SHA-256 integrity check of a downloaded snapshot. "
682            "Catches local disk corruption after extraction; relies on "
683            "manifest/file_checksums.json produced by `manifest`."
684        ),
685    )
686    p_vfy.add_argument(
687        "snapshot",
688        help="Snapshot root (the dir returned by snapshot_download or your staging dir).",
689    )
690    p_vfy.set_defaults(func=_cmd_verify)
691
692    p_lc = sub.add_parser(
693        "lossless-convert",
694        help="Re-encode rgb/ and fisheye/ PNGs to lossless WebP (~25% smaller, bit-identical pixels).",
695    )
696    p_lc.add_argument("--src", required=True, help="Source capture-tree root.")
697    p_lc.add_argument(
698        "--out",
699        required=True,
700        help="Output capture-tree root; must not overlap --src.",
701    )
702    p_lc.add_argument(
703        "--workers",
704        type=int,
705        default=max(1, (os.cpu_count() or 4) - 1),
706        help="Parallel encode workers (default: CPU count - 1).",
707    )
708    p_lc.add_argument(
709        "--dry-run",
710        action="store_true",
711        help="Walk the tree and report planned counts without writing anything.",
712    )
713    # Per-frame round-trip verification is mandatory from the CLI — it
714    # is the proof that every output frame is bit-faithful to its source.
715    # There is no --no-verify flag by design.
716    p_lc.add_argument(
717        "--overwrite-out",
718        action="store_true",
719        help="Allow writing into a non-empty --out directory.",
720    )
721    p_lc.add_argument(
722        "--webp-method",
723        type=int,
724        default=4,
725        choices=range(0, 7),
726        metavar="{0..6}",
727        help=(
728            "libwebp lossless effort level: 0=fastest/largest, 6=slowest/smallest. "
729            "Default 4: empirically the sweet spot on photo content (method=6 takes "
730            "10x longer for ~0%% extra savings on natural images)."
731        ),
732    )
733    p_lc.add_argument(
734        "--png-compress-level",
735        type=int,
736        default=9,
737        choices=range(0, 10),
738        metavar="{0..9}",
739        help=(
740            "PNG re-encode compression level (0-9). Default 9: densest, still "
741            "lossless. Applied to depth/* and sam2/masks/* PNGs."
742        ),
743    )
744    p_lc.add_argument(
745        "--skip-rgb-webp",
746        action="store_true",
747        help="Skip the rgb/fisheye/ego PNG → WebP-lossless pass.",
748    )
749    p_lc.add_argument(
750        "--skip-png-recompress",
751        action="store_true",
752        help="Skip the depth/masks PNG → PNG-level=9 re-encode pass.",
753    )
754    p_lc.add_argument(
755        "--skip-cam-pose",
756        action="store_true",
757        help="Skip the cam_pose .npz → per-frame .npy consolidation pass.",
758    )
759    p_lc.set_defaults(func=_cmd_lossless_convert)
760
761    p_cr = sub.add_parser("stage-croissant", help="Copy + patch the Croissant JSON.")
762    p_cr.add_argument(
763        "--staging", required=True, help="Staging dir; lands at <staging>/rpx_croissant.json."
764    )
765    p_cr.add_argument(
766        "--src",
767        default=None,
768        help="Source JSON path (default: <package>/dataset_hub/croissant_template.json).",
769    )
770    p_cr.add_argument(
771        "--repo-id",
772        default=DEFAULT_REPO_ID,
773        help=f"HF dataset repo id for url field (default: {DEFAULT_REPO_ID}).",
774    )
775    p_cr.add_argument("--version", default="1.0.0", help="Dataset version string (default: 1.0.0).")
776    p_cr.add_argument(
777        "--cite-as", default=None, help="Replace the citeAs field (e.g. a bibtex string)."
778    )
779    p_cr.add_argument(
780        "--overwrite", action="store_true", help="Overwrite an existing rpx_croissant.json."
781    )
782    p_cr.set_defaults(func=_cmd_stage_croissant)
783
784    return parser
785
786
787def main(argv: Optional[List[str]] = None) -> int:
788    from rpx_benchmark.cleanup import install_signal_cleanup
789
790    install_signal_cleanup()
791
792    args = build_parser().parse_args(argv)
793    return args.func(args)
794
795
796if __name__ == "__main__":
797    raise SystemExit(main())