rpx_benchmark.dataset_hub.dataset_card

  1"""Generate the HuggingFace dataset card (`README.md`) at the repo root.
  2
  3The dataset card is what HuggingFace renders on the dataset page and what
  4the Dataset Viewer parses for tags, task categories, and split sizes.
  5HF requires it as ``README.md`` at the repo root with YAML frontmatter.
  6
  7Inputs to the generator:
  8
  9* :class:`ScanResult` — for scene counts, modality file counts, byte totals.
 10* Optional :class:`PackResult` — for shard counts (more accurate than the
 11  scan when the packer collapses categories like sam2/* into masks_aux).
 12* Optional ``splits`` mapping — to populate the per-tier scene counts.
 13* Optional ``label_versions`` — to render the label versioning section.
 14
 15Output: ``<staging>/README.md``. Idempotent (overwrites with ``overwrite=True``).
 16"""
 17
 18from __future__ import annotations
 19
 20from dataclasses import dataclass
 21from pathlib import Path
 22from typing import Mapping, Optional
 23
 24from ..exceptions import DatasetError
 25from ..logging_utils import get_logger
 26from .recipes import (
 27    DEFAULT_REPO_ID,
 28    MULTI_OBJECT_TASK_RECIPES,
 29    SINGLE_OBJECT_TASK_RECIPES,
 30    SceneType,
 31)
 32from .scanner import ScanResult
 33
 34log = get_logger(__name__)
 35
 36
 37DEFAULT_LICENSE = "cc-by-4.0"
 38DEFAULT_PRETTY_NAME = "RPX: Robot Perception X"
 39
 40DEFAULT_TASK_CATEGORIES = (
 41    "image-segmentation",
 42    "depth-estimation",
 43    "object-detection",
 44    "visual-question-answering",
 45)
 46DEFAULT_TAGS = (
 47    "robotics",
 48    "embodied-ai",
 49    "rgb-d",
 50    "benchmark",
 51    "perception",
 52    "manipulation",
 53    "stereo",
 54)
 55
 56
 57@dataclass(frozen=True)
 58class CardSpec:
 59    """Knobs for the card generator."""
 60
 61    repo_id: str = DEFAULT_REPO_ID
 62    pretty_name: str = DEFAULT_PRETTY_NAME
 63    license: str = DEFAULT_LICENSE
 64    task_categories: tuple[str, ...] = DEFAULT_TASK_CATEGORIES
 65    tags: tuple[str, ...] = DEFAULT_TAGS
 66    paper_url: Optional[str] = None
 67    code_url: Optional[str] = "https://github.com/IRVLUTD/RPX"
 68
 69
 70def _human_bytes(n: int) -> str:
 71    f = float(n)
 72    for unit in ("B", "KB", "MB", "GB", "TB"):
 73        if f < 1024 or unit == "TB":
 74            return f"{f:.1f} {unit}" if unit != "B" else f"{int(f)} B"
 75        f /= 1024
 76    return f"{f:.1f} TB"
 77
 78
 79def _size_category(total_bytes: int) -> str:
 80    """HF size_categories tag bucket (used for search filters)."""
 81    gb = total_bytes / (1024**3)
 82    if gb < 1:
 83        return "100M<n<1B"
 84    if gb < 10:
 85        return "1B<n<10B"
 86    if gb < 100:
 87        return "10B<n<100B"
 88    if gb < 1024:
 89        return "100B<n<1T"
 90    return "n>1T"
 91
 92
 93def _yaml_list(values) -> str:
 94    return "\n".join(f"- {v}" for v in values)
 95
 96
 97def _frontmatter(spec: CardSpec, scan: ScanResult, splits: Mapping[str, str] | None = None) -> str:
 98    n_multi = len(scan.by_type(SceneType.MULTI_OBJECT))
 99    n_single = len(scan.by_type(SceneType.SINGLE_OBJECT))
100
101    lines = [
102        "---",
103        f"license: {spec.license}",
104        f'pretty_name: "{spec.pretty_name}"',
105        "task_categories:",
106        _yaml_list(spec.task_categories),
107        "language:",
108        "- en",
109        "tags:",
110        _yaml_list(spec.tags),
111        "size_categories:",
112        f"- {_size_category(scan.total_bytes)}",
113    ]
114
115    # If splits are provided, advertise them so HF Datasets viewer
116    # surfaces tier filters in the navigator.
117    if splits and n_multi > 0:
118        tier_counts: dict[str, int] = {"easy": 0, "medium": 0, "hard": 0}
119        for s in scan.by_type(SceneType.MULTI_OBJECT):
120            tier = splits.get(s.scene_id)
121            if tier in tier_counts:
122                tier_counts[tier] += 1
123        lines.append("configs:")
124        lines.append("- config_name: multi_object")
125        lines.append(
126            '  description: "MOS scenes (3 phases each), '
127            'tertile-cut by Effort-Stratified Difficulty."'
128        )
129        lines.append("  data_files:")
130        for tier, n in tier_counts.items():
131            if n:
132                lines.append(f'  - split: "{tier}"')
133                lines.append(f'    path: "splits/{tier}.txt"')
134        lines.append("- config_name: single_object")
135        lines.append(
136            '  description: "SOS scenes (one 360° collection per object); no difficulty split."'
137        )
138    lines.append("---")
139    return "\n".join(lines)
140
141
142def _body(
143    spec: CardSpec,
144    scan: ScanResult,
145    splits: Mapping[str, str] | None = None,
146    label_versions: Mapping[str, str] | None = None,
147) -> str:
148    n_multi = len(scan.by_type(SceneType.MULTI_OBJECT))
149    n_single = len(scan.by_type(SceneType.SINGLE_OBJECT))
150    modalities = scan.modality_totals()
151
152    mod_table_rows = []
153    for name, inv in modalities.items():
154        mod_table_rows.append(
155            f"| `{name}` | {inv.file_count:,} | {_human_bytes(inv.total_bytes)} |"
156        )
157    mod_table = "\n".join(mod_table_rows) if mod_table_rows else "*(none scanned)*"
158
159    multi_recipe_rows = "\n".join(
160        f"| `{name}` | {sorted(r.inputs)} → {sorted(r.labels)} |"
161        for name, r in MULTI_OBJECT_TASK_RECIPES.items()
162    )
163    single_recipe_rows = "\n".join(
164        f"| `{name}` | {sorted(r.inputs)} → {sorted(r.labels)} |"
165        for name, r in SINGLE_OBJECT_TASK_RECIPES.items()
166    )
167    label_v_rows = "\n".join(f"| `{k}` | `{v}` |" for k, v in (label_versions or {}).items())
168    label_v_block = (
169        "| modality | current version |\n|---|---|\n" + label_v_rows
170        if label_v_rows
171        else "*(default versions are in `manifest/current.json`)*"
172    )
173
174    paper_line = f"\n* Paper: {spec.paper_url}\n" if spec.paper_url else ""
175    code_line = f"* Code: {spec.code_url}\n" if spec.code_url else ""
176
177    return f"""
178# {spec.pretty_name}
179
180A real-world RGB-D benchmark for evaluating robot perception under
181embodied deployment conditions.
182
183{paper_line}{code_line}
184
185## Dataset at a glance
186
187| | |
188|---|---|
189| Multi-object scenes (MOS) | **{n_multi}** (3 phases each: clutter / interaction / clean) |
190| Single-object scenes (SOS) | **{n_single}** (one 360° collection per object) |
191| Total files | **{scan.file_count:,}** |
192| Total bytes | **{_human_bytes(scan.total_bytes)}** |
193
194## Modality inventory
195
196Three lossless space-saving steps were applied to this release; each is
197**bit-identical to the source** at the decoded pixel/value level
198(verified per-frame at conversion time):
199
200* `rgb/`, `fisheye/`, `ego/rgb/` are stored as **WebP-lossless**
201  (~45-55% smaller than the source PNGs on real-world photo content).
202* `depth/` and `sam2/masks/` are stored as **PNG re-encoded at
203  `compress_level=9`** (~20% smaller; same PNG file format, palette /
204  16-bit modes preserved exactly).
205* `cam_pose/` per-frame `.npz` files are stored as per-frame `.npy`
206  (~68% smaller; the per-file zip wrapper is removed and the position
207  and orientation are concatenated into one `(7,) float64` vector
208  packing `[x, y, z, qx, qy, qz, qw]`).
209
210Loaders see identical numpy arrays in every modality.
211
212| modality | files | bytes |
213|---|---:|---:|
214{mod_table}
215
216## Quick start
217
218```bash
219pip install "rpx-benchmark[hub]"
220hf auth login
221```
222
223```python
224from rpx_benchmark.dataset_hub import download_for_task
225
226# Pull just RGB + masks for the Easy difficulty tier — never the whole repo.
227res = download_for_task(task="segmentation", split="easy",
228                          repo_id="{spec.repo_id}")
229print(res.local_dir, res.matched_scenes)
230```
231
232```bash
233# Or from the CLI:
234python -m rpx_benchmark.dataset_hub.cli download \\
235    --task segmentation --split easy \\
236    --repo-id {spec.repo_id}
237```
238
239A subsequent call for a different task on the same split (e.g.
240`relative_pose`) reuses the cached RGB tars and only fetches the new
241modality (`cam_pose`) as the delta.
242
243## Repo layout
244
245```
246{spec.repo_id}/
247├── manifest/
248│   ├── frames_v1.parquet     # per-frame metadata (always pulled, ~30 MB)
249│   └── current.json          # default version per label modality
250├── splits/
251│   ├── scene_splits.json
252│   ├── easy.txt  medium.txt  hard.txt
253├── scenes/<scene_id>/<phase>/                     # MOS
254│   ├── rgb.tar  depth.tar  fisheye.tar
255│   └── labels/{{cam_pose,masks,masks_aux,sam2_meta,vqa}}/v1.tar
256├── objects/<object_id>/0/                         # SOS
257│   └── (same modality structure)
258├── objects_meta/                                  # questionnaire dedup
259│   ├── _index.json
260│   └── <object_id>/questionnaire.json
261└── README.md   ←  this file
262```
263
264## Tasks
265
266### Multi-object (use a difficulty split)
267
268| recipe | inputs → labels |
269|---|---|
270{multi_recipe_rows}
271
272### Single-object (no split — these are object templates)
273
274| recipe | inputs → labels |
275|---|---|
276{single_recipe_rows}
277
278## Label versioning
279
280Labels live at `labels/<name>/v<N>.tar`. Newer versions land at new
281paths; old versions stay reachable for reproducibility.
282
283{label_v_block}
284
285To pin to a specific version:
286
287```python
288download_for_task(
289    task="relative_pose", split="easy", repo_id="{spec.repo_id}",
290    label_versions={{"cam_pose": "v1"}},   # don't auto-upgrade to v2
291)
292```
293
294## Citation
295
296```bibtex
297@misc{{rpx2026,
298    title  = {{RPX: Robot Perception X — A real-world RGB-D benchmark for
299              embodied perception}},
300    author = {{IRVL UT Dallas}},
301    year   = 2026,
302    url    = {{https://huggingface.co/datasets/{spec.repo_id}}},
303}}
304```
305
306## License
307
308Released under the **{spec.license}** license.
309"""
310
311
312def write_dataset_card(
313    staging_root: Path,
314    scan: ScanResult,
315    spec: Optional[CardSpec] = None,
316    splits: Optional[Mapping[str, str]] = None,
317    label_versions: Optional[Mapping[str, str]] = None,
318    overwrite: bool = False,
319) -> Path:
320    """Render and write ``<staging_root>/README.md``.
321
322    Returns the destination path.
323    """
324    spec = spec or CardSpec()
325    out_path = Path(staging_root) / "README.md"
326    if out_path.exists() and not overwrite:
327        raise DatasetError(
328            f"refusing to overwrite existing dataset card: {out_path}",
329            hint="Pass overwrite=True or remove the file first.",
330        )
331
332    contents = (
333        _frontmatter(spec, scan, splits)
334        + "\n"
335        + _body(
336            spec,
337            scan,
338            splits,
339            label_versions,
340        )
341    )
342    out_path.parent.mkdir(parents=True, exist_ok=True)
343    out_path.write_text(contents, encoding="utf-8")
344    log.info("wrote dataset card → %s (%d bytes)", out_path, out_path.stat().st_size)
345    return out_path