rpx_benchmark.dataset_hub.downloader

  1"""Selective download from the RPX HuggingFace dataset repo.
  2
  3Resolves a ``(task, split, modalities, label_versions)`` request into a
  4list of glob patterns and hands them to ``snapshot_download``. The HF
  5cache is content-addressed by file hash, so re-pulling a previously
  6fetched modality is a no-op — that's how delta downloads work for free
  7when the team switches tasks on already-cached scenes.
  8
  9Wire-up::
 10
 11    pull = download_for_task(
 12        task="segmentation", split="easy", repo_id="IRVLUTD/RPX",
 13    )
 14    pull.local_dir       # local snapshot dir, ready to feed to a loader
 15    pull.manifest_table  # pyarrow.Table — per-frame metadata for the slice
 16    pull.bytes_fetched   # ≈ what the network actually moved (cache hits subtract)
 17"""
 18
 19from __future__ import annotations
 20
 21import json
 22from dataclasses import dataclass
 23from pathlib import Path
 24from typing import (
 25    TYPE_CHECKING,
 26    Any,
 27    Iterable,
 28    List,
 29    Mapping,
 30    Optional,
 31    Sequence,
 32    Set,
 33)
 34
 35from ..exceptions import ConfigError, DownloadError
 36from ..logging_utils import get_logger
 37from .manifest import SCHEMA_VERSION
 38from .packer import SCENE_ROOT_BY_TYPE
 39from .recipes import (
 40    DEFAULT_REPO_ID,
 41    QUESTIONNAIRE,
 42    SceneType,
 43    TaskRecipe,
 44    resolve_recipe,
 45)
 46
 47#: Scene types whose scenes carry a real difficulty split (ego inherits its
 48#: split from its mos sibling via a shared scene_id — see ego_layout.py).
 49_SPLIT_BEARING_TYPES = (SceneType.MULTI_OBJECT, SceneType.EGO)
 50
 51if TYPE_CHECKING:  # pragma: no cover - only for type hints
 52    import pyarrow as pa
 53
 54
 55log = get_logger(__name__)
 56
 57
 58# Modalities the packer treats as raw — no version suffix in the repo path.
 59# Must stay in sync with packer.RAW_MODALITIES.
 60_RAW_MODALITIES = frozenset({"rgb", "depth", "fisheye"})
 61
 62# Modalities that live OUTSIDE the per-(scene, phase) tar shards, under
 63# the shared ``objects_meta/<object_id>/`` directory. Their pull patterns
 64# are derived from the per-scene mask_to_object lookups, not from
 65# scene-rooted globs.
 66_SHARED_MODALITIES = frozenset({QUESTIONNAIRE})
 67
 68
 69@dataclass(frozen=True)
 70class DownloadResult:
 71    """Outcome of one selective pull."""
 72
 73    repo_id: str
 74    revision: Optional[str]
 75    task: str
 76    split: Optional[str]
 77    scene_type: SceneType
 78    local_dir: Path
 79    allow_patterns: List[str]
 80    matched_scenes: List[str]
 81    bytes_fetched: int
 82    files_fetched: int
 83    cache_hits: int = 0
 84    manifest_table: Optional["pa.Table"] = None
 85
 86
 87def _hub():
 88    try:
 89        import huggingface_hub as hf  # noqa: WPS433
 90    except ImportError as e:
 91        raise DownloadError(
 92            "huggingface_hub is required for RPX downloads.",
 93            hint="Install with: pip install 'rpx-benchmark[hub]'",
 94        ) from e
 95    return hf
 96
 97
 98def _arrow():
 99    try:
100        import pyarrow as pa  # noqa: WPS433
101        import pyarrow.parquet as pq  # noqa: WPS433
102    except ImportError as e:
103        raise DownloadError(
104            "pyarrow is required to read the dataset manifest.",
105            hint="Install with: pip install 'rpx-benchmark[hub]'",
106        ) from e
107    return pa, pq
108
109
110# --------------------------------------------------------------------- #
111# Manifest helpers (pulled and read first; small file, always cheap)
112# --------------------------------------------------------------------- #
113
114
115def _fetch_manifest(
116    hf,
117    repo_id: str,
118    revision: Optional[str],
119    cache_dir: Optional[Path],
120) -> tuple[Path, Path]:
121    """Download just ``manifest/{frames_v1.parquet, current.json}``."""
122    parquet_path = hf.hf_hub_download(
123        repo_id=repo_id,
124        repo_type="dataset",
125        filename=f"manifest/frames_{SCHEMA_VERSION}.parquet",
126        revision=revision,
127        cache_dir=str(cache_dir) if cache_dir else None,
128    )
129    current_path = hf.hf_hub_download(
130        repo_id=repo_id,
131        repo_type="dataset",
132        filename="manifest/current.json",
133        revision=revision,
134        cache_dir=str(cache_dir) if cache_dir else None,
135    )
136    return Path(parquet_path), Path(current_path)
137
138
139def _filter_manifest(
140    table: "pa.Table",
141    scene_type: SceneType,
142    split: Optional[str],
143) -> "pa.Table":
144    pa, _ = _arrow()
145    import pyarrow.compute as pc  # noqa: WPS433
146
147    mask = pc.equal(table["scene_type"], pa.scalar(scene_type.value))
148    if split is not None and scene_type in _SPLIT_BEARING_TYPES:
149        mask = pc.and_(mask, pc.equal(table["split"], pa.scalar(split)))
150    return table.filter(mask)
151
152
153# --------------------------------------------------------------------- #
154# Pattern resolution
155# --------------------------------------------------------------------- #
156
157
158def _resolve_label_version(
159    modality: str,
160    current: Mapping[str, Any],
161    explicit: Mapping[str, str],
162) -> str:
163    """Pick the label version for ``modality``: explicit > current.json > 'v1'."""
164    if modality in explicit:
165        return explicit[modality]
166    versions = current.get("label_versions", {})
167    return versions.get(modality, "v1")
168
169
170def _scene_root_for(scene_type: SceneType) -> str:
171    return SCENE_ROOT_BY_TYPE[scene_type]
172
173
174def _phase_glob_segment(scene_type: SceneType) -> str:
175    """The glob segment standing in for "phase" in an allow_pattern.
176
177    ego now nests under the SAME scenes/<scene>/ root as mos (as a
178    literal "ego" segment alongside the numeric 0/1/2/... phases -- see
179    packer.phase_segment), so a bare "*" would match BOTH: a mos download
180    would incidentally also sweep up scenes/<scene>/ego/rgb.tar, and vice
181    versa. "[0-9]*" matches only digit-leading (numeric-phase) segments;
182    ego's literal "ego" doesn't match it, and "ego" is used verbatim for
183    ego scene_type -- so each family's pull stays scoped to itself.
184    """
185    return "ego" if scene_type is SceneType.EGO else "[0-9]*"
186
187
188def _build_allow_patterns(
189    matched_scenes: Iterable[str],
190    scene_type: SceneType,
191    modalities: Iterable[str],
192    current: Mapping[str, Any],
193    label_versions: Mapping[str, str],
194) -> List[str]:
195    """One glob per (scene, modality). Shared modalities (questionnaire)
196    are *not* expanded here — see :func:`_shared_artefact_patterns` for
197    the per-object dedup pull.
198    """
199    root = _scene_root_for(scene_type)
200    phase_seg = _phase_glob_segment(scene_type)
201    patterns: List[str] = []
202    for scene in sorted(matched_scenes):
203        for m in sorted(modalities):
204            if m in _SHARED_MODALITIES:
205                continue  # handled by _shared_artefact_patterns
206            if m in _RAW_MODALITIES:
207                patterns.append(f"{root}/{scene}/{phase_seg}/{m}.tar")
208            else:
209                v = _resolve_label_version(m, current, label_versions)
210                patterns.append(f"{root}/{scene}/{phase_seg}/labels/{m}/{v}.tar")
211    return patterns
212
213
214def _shared_artefact_patterns(modalities: Iterable[str]) -> List[str]:
215    """Build allow_patterns for the ``objects_meta/`` layer.
216
217    For now we pull *all* per-object questionnaires when a recipe asks
218    for ``questionnaire`` — they are tiny (~220 small JSON files). A
219    future optimisation would resolve the exact set of object IDs from
220    the matched scenes' ``mask_to_object.json`` files (already in the
221    sam2_meta tars), but the saving is negligible.
222    """
223    patterns: List[str] = []
224    if QUESTIONNAIRE in modalities:
225        patterns.append("objects_meta/*/questionnaire.json")
226    return patterns
227
228
229# --------------------------------------------------------------------- #
230# Public API
231# --------------------------------------------------------------------- #
232
233
234def download_for_task(
235    task: str,
236    split: Optional[str] = None,
237    repo_id: str = DEFAULT_REPO_ID,
238    *,
239    revision: Optional[str] = None,
240    cache_dir: Optional[Path] = None,
241    extra_modalities: Sequence[str] = (),
242    label_versions: Optional[Mapping[str, str]] = None,
243    scene_type: Optional[SceneType] = None,
244    max_workers: int = 8,
245    return_manifest: bool = True,
246) -> DownloadResult:
247    """Pull just the tar shards needed to run ``task`` on ``split``.
248
249    Parameters
250    ----------
251    task : str
252        Recipe key from :data:`MULTI_OBJECT_TASK_RECIPES` or
253        :data:`SINGLE_OBJECT_TASK_RECIPES`.
254    split : str, optional
255        ``"easy" | "medium" | "hard"``. Required for multi-object tasks,
256        forbidden for single-object tasks (which have no splits).
257    repo_id : str
258        HuggingFace dataset repo (default: ``IRVLUTD/RPX``).
259    extra_modalities : Sequence[str]
260        Add modalities beyond the recipe's defaults (e.g. ``["masks_aux"]``
261        if you want the visualisation overlays alongside ``masks``).
262    label_versions : Mapping[str, str], optional
263        Override which label version to pull for one or more modalities.
264        Default behaviour is to follow the repo's ``manifest/current.json``.
265    scene_type : SceneType, optional
266        Disambiguates when a task name exists in both recipe tables.
267
268    Returns
269    -------
270    DownloadResult
271        ``local_dir`` is the local snapshot root; combine with
272        ``allow_patterns`` and ``manifest_table`` to feed a downstream
273        loader. ``bytes_fetched`` reports the size on disk after the pull
274        (for true delta accounting prefer your own ``du -sh``).
275    """
276    hf = _hub()
277    recipe: TaskRecipe = resolve_recipe(task, scene_type=scene_type)
278
279    if recipe.scene_type is SceneType.SINGLE_OBJECT and split is not None:
280        raise ConfigError(
281            f"task {task!r} targets single-object scenes, which have no splits.",
282            hint="Drop the --split argument for single-object recipes.",
283        )
284    if recipe.scene_type in _SPLIT_BEARING_TYPES and split is None:
285        raise ConfigError(
286            f"task {task!r} targets {recipe.scene_type.value} scenes — pick a split.",
287            hint="Pass split='easy' | 'medium' | 'hard'.",
288        )
289
290    parquet_path, current_path = _fetch_manifest(
291        hf,
292        repo_id,
293        revision,
294        cache_dir,
295    )
296    current = json.loads(current_path.read_text(encoding="utf-8"))
297    _, pq = _arrow()
298    table = pq.read_table(parquet_path)
299    sliced = _filter_manifest(table, recipe.scene_type, split)
300
301    matched_scenes: Set[str] = set(sliced["scene_id"].to_pylist())
302    if not matched_scenes:
303        raise DownloadError(
304            f"No scenes matched task={task!r} split={split!r} on {repo_id}.",
305            hint=(
306                "Verify the manifest has rows for this scene_type/split. "
307                "If the dataset was uploaded without a `splits=...` argument "
308                "to build_frame_manifest, multi-object rows will all have "
309                "split=null."
310            ),
311        )
312
313    modalities = set(recipe.all_modalities()) | set(extra_modalities)
314    explicit = dict(label_versions or {})
315    patterns = _build_allow_patterns(
316        matched_scenes,
317        recipe.scene_type,
318        modalities,
319        current,
320        explicit,
321    )
322    patterns.extend(_shared_artefact_patterns(modalities))
323
324    log.info(
325        "download plan: task=%s split=%s scenes=%d modalities=%s patterns=%d",
326        task,
327        split,
328        len(matched_scenes),
329        sorted(modalities),
330        len(patterns),
331    )
332
333    try:
334        local_dir = hf.snapshot_download(
335            repo_id=repo_id,
336            repo_type="dataset",
337            allow_patterns=patterns,
338            revision=revision,
339            cache_dir=str(cache_dir) if cache_dir else None,
340            max_workers=max_workers,
341        )
342    except Exception as e:
343        raise DownloadError(
344            f"snapshot_download failed for {repo_id}: {e}",
345            hint=(
346                "Re-run to retry; HF caches partial downloads. If auth is the "
347                "issue, run `hf auth login` and confirm read access."
348            ),
349        ) from e
350
351    local = Path(local_dir)
352
353    # Mirror the post-fetch state that hub.download_split produces: extract
354    # every tar shard to flat per-modality dirs so downstream consumers
355    # (loader, viz, custom scripts) hit pure disk reads instead of opening
356    # and seeking tars per frame. Idempotent — files already on disk skip.
357    from ..hub import _extract_snapshot_tars
358    _extract_snapshot_tars(local)
359
360    files = [p for p in local.rglob("*") if p.is_file()]
361    bytes_ = sum(p.stat().st_size for p in files)
362
363    return DownloadResult(
364        repo_id=repo_id,
365        revision=revision,
366        task=task,
367        split=split,
368        scene_type=recipe.scene_type,
369        local_dir=local,
370        allow_patterns=patterns,
371        matched_scenes=sorted(matched_scenes),
372        bytes_fetched=bytes_,
373        files_fetched=len(files),
374        cache_hits=0,  # HF doesn't expose per-call hit count
375        manifest_table=sliced if return_manifest else None,
376    )