rpx_benchmark.dataset_hub
RPX dataset hub — selective HuggingFace download/upload pipeline.
Layered on top of rpx_benchmark.hub (which already implements
modality-aware snapshot_download) and adds:
rpx_benchmark.dataset_hub.recipes— task → modality recipes (multi-object + single-object).rpx_benchmark.dataset_hub.mock— synthetic capture-tree generator for tests / dry runs.rpx_benchmark.dataset_hub.scanner— walks the on-disk capture layout, reports inventory.rpx_benchmark.dataset_hub.packer— converts a capture tree into per-modality tar shards.rpx_benchmark.dataset_hub.manifest— writes the per-frame Parquet manifest the hub ships.rpx_benchmark.dataset_hub.uploader— pushes a packed staging dir to a HF dataset repo.rpx_benchmark.dataset_hub.cli—python -m rpx_benchmark.dataset_hub.clientry points.
End-to-end usage::
src = Path("/path/to/test_dataset_aggregated")
stage = Path("/tmp/rpx_staging")
scan = scan_capture_root(src)
pack = pack_capture_tree(PackPlan(src, stage), scan)
build_frame_manifest(scan, pack, stage, splits=splits_dict)
upload_staging(UploadPlan(stage, repo_id="IRVLUTD/RPX"))
1"""RPX dataset hub — selective HuggingFace download/upload pipeline. 2 3Layered on top of :mod:`rpx_benchmark.hub` (which already implements 4modality-aware ``snapshot_download``) and adds: 5 6 * :mod:`.recipes` — task → modality recipes (multi-object + single-object). 7 * :mod:`.mock` — synthetic capture-tree generator for tests / dry runs. 8 * :mod:`.scanner` — walks the on-disk capture layout, reports inventory. 9 * :mod:`.packer` — converts a capture tree into per-modality tar shards. 10 * :mod:`.manifest` — writes the per-frame Parquet manifest the hub ships. 11 * :mod:`.uploader` — pushes a packed staging dir to a HF dataset repo. 12 * :mod:`.cli` — ``python -m rpx_benchmark.dataset_hub.cli`` entry points. 13 14End-to-end usage:: 15 16 src = Path("/path/to/test_dataset_aggregated") 17 stage = Path("/tmp/rpx_staging") 18 scan = scan_capture_root(src) 19 pack = pack_capture_tree(PackPlan(src, stage), scan) 20 build_frame_manifest(scan, pack, stage, splits=splits_dict) 21 upload_staging(UploadPlan(stage, repo_id="IRVLUTD/RPX")) 22""" 23 24from __future__ import annotations 25 26from .croissant import ( 27 CroissantPatch, 28 stage_croissant, 29) 30from .dataset_card import ( 31 CardSpec, 32 write_dataset_card, 33) 34from .downloader import ( 35 DownloadResult, 36 download_for_task, 37) 38from .manifest import ( 39 SCHEMA_VERSION, 40 ManifestPaths, 41 build_frame_manifest, 42 read_frame_manifest, 43) 44from .packer import ( 45 RAW_MODALITIES, 46 PackedShard, 47 PackPlan, 48 PackResult, 49 SharedArtefact, 50 pack_capture_tree, 51 pack_objects_meta, 52) 53from .recipes import ( 54 DEFAULT_REPO_ID, 55 MULTI_OBJECT_TASK_RECIPES, 56 SINGLE_OBJECT_TASK_RECIPES, 57 SceneType, 58 TaskRecipe, 59 resolve_recipe, 60) 61from .scanner import ( 62 ModalityInventory, 63 PhaseInventory, 64 ScanResult, 65 SceneInventory, 66 scan_capture_root, 67) 68from .split_manifests import ( 69 write_split_manifests, 70) 71from .staging import ( 72 SCENE_SPLITS_JSON, 73 SPLIT_TIER_FILES, 74 StagedFile, 75 load_scene_splits, 76 stage_paths, 77 stage_splits, 78) 79from .uploader import ( 80 DEFAULT_IGNORE_PATTERNS, 81 UploadPlan, 82 UploadResult, 83 upload_paths, 84 upload_staging, 85) 86 87__all__ = [ 88 # recipes 89 "DEFAULT_REPO_ID", 90 "MULTI_OBJECT_TASK_RECIPES", 91 "SINGLE_OBJECT_TASK_RECIPES", 92 "SceneType", 93 "TaskRecipe", 94 "resolve_recipe", 95 # scanner 96 "ModalityInventory", 97 "PhaseInventory", 98 "SceneInventory", 99 "ScanResult", 100 "scan_capture_root", 101 # packer 102 "RAW_MODALITIES", 103 "PackPlan", 104 "PackResult", 105 "PackedShard", 106 "SharedArtefact", 107 "pack_capture_tree", 108 "pack_objects_meta", 109 # manifest 110 "SCHEMA_VERSION", 111 "ManifestPaths", 112 "build_frame_manifest", 113 "read_frame_manifest", 114 # split manifests (per-task, per-split) 115 "write_split_manifests", 116 # uploader 117 "DEFAULT_IGNORE_PATTERNS", 118 "UploadPlan", 119 "UploadResult", 120 "upload_paths", 121 "upload_staging", 122 # croissant 123 "CroissantPatch", 124 "stage_croissant", 125 # dataset card 126 "CardSpec", 127 "write_dataset_card", 128 # downloader 129 "DownloadResult", 130 "download_for_task", 131 # staging 132 "SPLIT_TIER_FILES", 133 "SCENE_SPLITS_JSON", 134 "StagedFile", 135 "load_scene_splits", 136 "stage_paths", 137 "stage_splits", 138]
DEFAULT_REPO_ID =
$GITHUB_REPOSITORY
MULTI_OBJECT_TASK_RECIPES =
{'monocular_depth': TaskRecipe(name='monocular_depth', scene_type=<SceneType.MULTI_OBJECT: 'multi_object'>, inputs=frozenset({'rgb'}), labels=frozenset({'depth'}), notes='Paper task Image Depth. Frame-level: model receives one RGB frame, outputs one depth map. Depth captured by D435 is the GT.'), 'video_depth': TaskRecipe(name='video_depth', scene_type=<SceneType.MULTI_OBJECT: 'multi_object'>, inputs=frozenset({'rgb'}), labels=frozenset({'depth'}), notes="Paper task Video Depth. Video-level: model receives the full ~250-frame phase clip as RGB sequence, outputs a per-frame depth sequence. Same scenes and GT as monocular_depth; differs in iteration unit and metric set (per-frame metrics + temporal metrics — TODO: finalise from Feynman's research in benchmark/README.md)."), 'rgbd_segmentation': TaskRecipe(name='rgbd_segmentation', scene_type=<SceneType.MULTI_OBJECT: 'multi_object'>, inputs=frozenset({'depth', 'rgb'}), labels=frozenset({'masks'}), notes=''), 'segmentation': TaskRecipe(name='segmentation', scene_type=<SceneType.MULTI_OBJECT: 'multi_object'>, inputs=frozenset({'rgb'}), labels=frozenset({'masks'}), notes=''), 'relative_pose': TaskRecipe(name='relative_pose', scene_type=<SceneType.MULTI_OBJECT: 'multi_object'>, inputs=frozenset({'rgb'}), labels=frozenset({'cam_pose'}), notes=''), 'rgbd_relative_pose': TaskRecipe(name='rgbd_relative_pose', scene_type=<SceneType.MULTI_OBJECT: 'multi_object'>, inputs=frozenset({'depth', 'rgb'}), labels=frozenset({'cam_pose'}), notes=''), 'stereo_depth': TaskRecipe(name='stereo_depth', scene_type=<SceneType.MULTI_OBJECT: 'multi_object'>, inputs=frozenset({'fisheye'}), labels=frozenset({'depth'}), notes='T265 stereo fisheye → depth. D435 depth is the GT after rectification.'), 'object_tracking': TaskRecipe(name='object_tracking', scene_type=<SceneType.MULTI_OBJECT: 'multi_object'>, inputs=frozenset({'rgb'}), labels=frozenset({'masks'}), notes=''), 'vqa': TaskRecipe(name='vqa', scene_type=<SceneType.MULTI_OBJECT: 'multi_object'>, inputs=frozenset({'rgb'}), labels=frozenset({'questionnaire', 'vqa'}), notes='Reserved slot — VQA labels land in v1.x.')}
SINGLE_OBJECT_TASK_RECIPES =
{'object_templates': TaskRecipe(name='object_templates', scene_type=<SceneType.SINGLE_OBJECT: 'single_object'>, inputs=frozenset({'rgb'}), labels=frozenset({'masks'}), notes='360° turntable RGB + mask — feed as reference views for in-context tasks.'), 'object_templates_rgbd': TaskRecipe(name='object_templates_rgbd', scene_type=<SceneType.SINGLE_OBJECT: 'single_object'>, inputs=frozenset({'depth', 'rgb'}), labels=frozenset({'masks'}), notes=''), 'object_pose_library': TaskRecipe(name='object_pose_library', scene_type=<SceneType.SINGLE_OBJECT: 'single_object'>, inputs=frozenset({'depth', 'rgb'}), labels=frozenset({'masks', 'cam_pose'}), notes='Per-object pose graph for retrieval / pose initialisation.')}
class
SceneType(builtins.str, enum.Enum):
32class SceneType(str, Enum): 33 """Top-level partition: which scene family a recipe targets.""" 34 35 MULTI_OBJECT = "multi_object" 36 SINGLE_OBJECT = "single_object" 37 EGO = "ego"
Top-level partition: which scene family a recipe targets.
MULTI_OBJECT =
<SceneType.MULTI_OBJECT: 'multi_object'>
SINGLE_OBJECT =
<SceneType.SINGLE_OBJECT: 'single_object'>
EGO =
<SceneType.EGO: 'ego'>
@dataclass(frozen=True)
class
TaskRecipe:
40@dataclass(frozen=True) 41class TaskRecipe: 42 """Declarative spec: 'task X needs these modalities for inputs and labels'. 43 44 The split between ``inputs`` and ``labels`` is purely a documentation 45 aid — the downloader unions them when expanding to glob patterns. 46 Storing them separately makes it obvious *why* each modality is in the 47 recipe (which matters when modalities are added or renamed). 48 """ 49 50 name: str 51 scene_type: SceneType 52 inputs: FrozenSet[str] 53 labels: FrozenSet[str] = field(default_factory=frozenset) 54 notes: str = "" 55 56 def all_modalities(self) -> FrozenSet[str]: 57 return self.inputs | self.labels
Declarative spec: 'task X needs these modalities for inputs and labels'.
The split between inputs and labels is purely a documentation
aid — the downloader unions them when expanding to glob patterns.
Storing them separately makes it obvious why each modality is in the
recipe (which matters when modalities are added or renamed).
TaskRecipe( name: str, scene_type: SceneType, inputs: FrozenSet[str], labels: FrozenSet[str] = <factory>, notes: str = '')
scene_type: SceneType
218def resolve_recipe( 219 task: str, 220 scene_type: Optional[SceneType] = None, 221) -> TaskRecipe: 222 """Look up a recipe by ``task`` name. 223 224 If ``scene_type`` is given, only that table is searched (use this when a 225 name collides between multi- and single-object tables). Otherwise we 226 search multi-object first, then single-object, and raise on miss. 227 """ 228 if scene_type is SceneType.MULTI_OBJECT: 229 return MULTI_OBJECT_TASK_RECIPES[task] 230 if scene_type is SceneType.SINGLE_OBJECT: 231 return SINGLE_OBJECT_TASK_RECIPES[task] 232 if scene_type is SceneType.EGO: 233 return EGO_TASK_RECIPES[task] 234 if task in MULTI_OBJECT_TASK_RECIPES: 235 return MULTI_OBJECT_TASK_RECIPES[task] 236 if task in SINGLE_OBJECT_TASK_RECIPES: 237 return SINGLE_OBJECT_TASK_RECIPES[task] 238 if task in EGO_TASK_RECIPES: 239 return EGO_TASK_RECIPES[task] 240 raise ConfigError( 241 f"Unknown task {task!r}.", 242 hint=( 243 "Known multi-object recipes: " 244 f"{sorted(MULTI_OBJECT_TASK_RECIPES)}. " 245 "Known single-object recipes: " 246 f"{sorted(SINGLE_OBJECT_TASK_RECIPES)}. " 247 "Known ego recipes: " 248 f"{sorted(EGO_TASK_RECIPES)}. " 249 "Pass `scene_type=SceneType.SINGLE_OBJECT` if a name collides " 250 "between tables." 251 ), 252 )
Look up a recipe by task name.
If scene_type is given, only that table is searched (use this when a
name collides between multi- and single-object tables). Otherwise we
search multi-object first, then single-object, and raise on miss.
@dataclass(frozen=True)
class
ModalityInventory:
53@dataclass(frozen=True) 54class ModalityInventory: 55 """Inventory of a single modality directory under one phase.""" 56 57 name: str # e.g. "rgb", "depth", "fisheye", "cam_pose", "sam2/masks" 58 file_count: int 59 total_bytes: int
Inventory of a single modality directory under one phase.
@dataclass(frozen=True)
class
PhaseInventory:
62@dataclass(frozen=True) 63class PhaseInventory: 64 """Inventory of one phase directory (set of modalities).""" 65 66 phase_index: int 67 modalities: Dict[str, ModalityInventory] 68 69 @property 70 def total_bytes(self) -> int: 71 return sum(m.total_bytes for m in self.modalities.values()) 72 73 @property 74 def file_count(self) -> int: 75 return sum(m.file_count for m in self.modalities.values())
Inventory of one phase directory (set of modalities).
PhaseInventory( phase_index: int, modalities: Dict[str, ModalityInventory])
modalities: Dict[str, ModalityInventory]
@dataclass(frozen=True)
class
SceneInventory:
78@dataclass(frozen=True) 79class SceneInventory: 80 """Inventory of one scene directory (set of phases).""" 81 82 scene_id: str 83 scene_type: SceneType 84 phases: List[PhaseInventory] 85 86 @property 87 def total_bytes(self) -> int: 88 return sum(p.total_bytes for p in self.phases) 89 90 @property 91 def file_count(self) -> int: 92 return sum(p.file_count for p in self.phases)
Inventory of one scene directory (set of phases).
SceneInventory( scene_id: str, scene_type: SceneType, phases: List[PhaseInventory])
scene_type: SceneType
phases: List[PhaseInventory]
@dataclass(frozen=True)
class
ScanResult:
95@dataclass(frozen=True) 96class ScanResult: 97 """Top-level scan result over a capture root.""" 98 99 root: Path 100 scenes: List[SceneInventory] 101 skipped: List[Path] = field(default_factory=list) 102 103 # ── aggregate helpers ──────────────────────────────────────────── 104 def by_type(self, scene_type: SceneType) -> List[SceneInventory]: 105 return [s for s in self.scenes if s.scene_type is scene_type] 106 107 @property 108 def total_bytes(self) -> int: 109 return sum(s.total_bytes for s in self.scenes) 110 111 @property 112 def file_count(self) -> int: 113 return sum(s.file_count for s in self.scenes) 114 115 def modality_totals(self) -> Dict[str, ModalityInventory]: 116 """Aggregate per-modality file counts and bytes across all phases.""" 117 agg: Dict[str, list[int]] = {} 118 for scene in self.scenes: 119 for phase in scene.phases: 120 for name, inv in phase.modalities.items(): 121 if name not in agg: 122 agg[name] = [0, 0] 123 agg[name][0] += inv.file_count 124 agg[name][1] += inv.total_bytes 125 return { 126 name: ModalityInventory(name=name, file_count=fc, total_bytes=tb) 127 for name, (fc, tb) in sorted(agg.items()) 128 }
Top-level scan result over a capture root.
ScanResult( root: pathlib.Path, scenes: List[SceneInventory], skipped: List[pathlib.Path] = <factory>)
scenes: List[SceneInventory]
115 def modality_totals(self) -> Dict[str, ModalityInventory]: 116 """Aggregate per-modality file counts and bytes across all phases.""" 117 agg: Dict[str, list[int]] = {} 118 for scene in self.scenes: 119 for phase in scene.phases: 120 for name, inv in phase.modalities.items(): 121 if name not in agg: 122 agg[name] = [0, 0] 123 agg[name][0] += inv.file_count 124 agg[name][1] += inv.total_bytes 125 return { 126 name: ModalityInventory(name=name, file_count=fc, total_bytes=tb) 127 for name, (fc, tb) in sorted(agg.items()) 128 }
Aggregate per-modality file counts and bytes across all phases.
194def scan_capture_root(root: Path) -> ScanResult: 195 """Scan a capture-tree root and return a :class:`ScanResult`. 196 197 Walks ``root/mos/*`` (multi-object scenes) and ``root/sos/*`` (single- 198 object scenes). Anything else under root, and anything inside ``mos/`` 199 or ``sos/`` that isn't a scene directory, is collected into 200 :attr:`ScanResult.skipped` so the caller can flag typos / strays. 201 """ 202 root = Path(root) 203 if not root.is_dir(): 204 raise FileNotFoundError(f"capture root does not exist: {root}") 205 206 scenes: List[SceneInventory] = [] 207 skipped: List[Path] = [] 208 for entry in sorted(root.iterdir()): 209 # Top-level: only ``mos/`` and ``sos/`` are recognised. 210 if not entry.is_dir() or entry.name not in _TYPE_DIRS: 211 skipped.append(entry) 212 continue 213 214 scene_type = _TYPE_DIRS[entry.name] 215 for scene_dir in sorted(entry.iterdir()): 216 if not scene_dir.is_dir(): 217 skipped.append(scene_dir) 218 continue 219 phase_dirs = sorted( 220 [p for p in scene_dir.iterdir() if p.is_dir()], 221 key=lambda p: p.name, 222 ) 223 phases = [_scan_phase(p) for p in phase_dirs] 224 scenes.append( 225 SceneInventory( 226 scene_id=scene_dir.name, 227 scene_type=scene_type, 228 phases=phases, 229 ) 230 ) 231 232 return ScanResult(root=root, scenes=scenes, skipped=skipped)
Scan a capture-tree root and return a ScanResult.
Walks root/mos/* (multi-object scenes) and root/sos/* (single-
object scenes). Anything else under root, and anything inside mos/
or sos/ that isn't a scene directory, is collected into
ScanResult.skipped so the caller can flag typos / strays.
RAW_MODALITIES =
frozenset({'depth', 'fisheye', 'rgb'})
@dataclass(frozen=True)
class
PackPlan:
105@dataclass(frozen=True) 106class PackPlan: 107 """Knobs for the packer.""" 108 109 src_root: Path 110 staging_root: Path 111 label_version: str = "v1" 112 overwrite: bool = False 113 compute_hash: bool = True
Knobs for the packer.
@dataclass(frozen=True)
class
PackResult:
134@dataclass(frozen=True) 135class PackResult: 136 """Aggregate result of one packer run.""" 137 138 shards: List[PackedShard] 139 skipped_keys: List[str] = field(default_factory=list) 140 141 @property 142 def total_bytes(self) -> int: 143 return sum(s.total_bytes for s in self.shards) 144 145 @property 146 def total_files(self) -> int: 147 return sum(s.file_count for s in self.shards)
Aggregate result of one packer run.
PackResult( shards: List[PackedShard], skipped_keys: List[str] = <factory>)
shards: List[PackedShard]
@dataclass(frozen=True)
class
PackedShard:
116@dataclass(frozen=True) 117class PackedShard: 118 """One tar shard the packer produced.""" 119 120 repo_path: str # path relative to staging_root, forward-slashes 121 scene_id: str 122 scene_type: SceneType # disambiguates scene_id collisions across families 123 # (ego scenes deliberately reuse their mos/ 124 # sibling's scene_id — see manifest.py's 125 # _index_shards_by_phase / _frame_filenames_for) 126 phase: int 127 modality: str 128 is_label: bool 129 file_count: int 130 total_bytes: int 131 sha256: Optional[str] = None
One tar shard the packer produced.
PackedShard( repo_path: str, scene_id: str, scene_type: SceneType, phase: int, modality: str, is_label: bool, file_count: int, total_bytes: int, sha256: Optional[str] = None)
scene_type: SceneType
288def pack_capture_tree(plan: PackPlan, scan: ScanResult) -> PackResult: 289 """Pack each ``(scene, phase, modality)`` triple in ``scan`` into one tar. 290 291 Reusing :class:`ScanResult` from the scanner avoids re-walking the 292 source tree — the scan already has the file counts, byte totals, and 293 structure we need. The packer just translates each (scene, phase, 294 modality) entry into a tar at the right repo path. 295 """ 296 if plan.src_root != scan.root: 297 raise ConfigError( 298 f"plan.src_root ({plan.src_root}) does not match the root the " 299 f"scan was taken at ({scan.root}).", 300 hint="Pass the same root to scan_capture_root() and PackPlan.", 301 ) 302 plan.staging_root.mkdir(parents=True, exist_ok=True) 303 304 shards: List[PackedShard] = [] 305 skipped: List[str] = [] 306 307 for scene in scan.scenes: 308 for phase in scene.phases: 309 # Group source modalities by output (modality, is_label) so 310 # masks_aux/* all flow into the same masks_aux.tar. 311 grouped: Dict[tuple[str, bool], list[str]] = {} 312 for scanner_key in phase.modalities: 313 target = _MODALITY_PLAN.get(scanner_key) 314 if target is None: 315 skipped.append(f"{scene.scene_id}/{phase.phase_index}/{scanner_key}") 316 continue 317 grouped.setdefault(target, []).append(scanner_key) 318 319 for (out_modality, is_label), keys in grouped.items(): 320 files: List[Path] = [] 321 for k in keys: 322 files.extend( 323 _enumerate_files_for_modality( 324 plan.src_root, 325 scene.scene_type, 326 scene.scene_id, 327 phase.phase_index, 328 k, 329 ) 330 ) 331 if not files: 332 continue 333 334 repo_path = _shard_repo_path( 335 scene.scene_type, 336 scene.scene_id, 337 phase.phase_index, 338 out_modality, 339 is_label, 340 plan.label_version, 341 ) 342 out_path = plan.staging_root / repo_path 343 if out_path.exists() and not plan.overwrite: 344 raise DatasetError( 345 f"refusing to overwrite existing shard: {out_path}", 346 hint="Pass overwrite=True or remove the staging dir.", 347 ) 348 349 fc, bs, sha = _write_tar( 350 out_path, 351 files, 352 rel_to=( 353 _scene_src_root(plan.src_root, scene.scene_type, scene.scene_id) 354 / str(phase.phase_index) 355 ), 356 compute_hash=plan.compute_hash, 357 ) 358 shards.append( 359 PackedShard( 360 repo_path=repo_path, 361 scene_id=scene.scene_id, 362 scene_type=scene.scene_type, 363 phase=phase.phase_index, 364 modality=out_modality, 365 is_label=is_label, 366 file_count=fc, 367 total_bytes=bs, 368 sha256=sha, 369 ) 370 ) 371 log.info("packed %s (%d files, %d bytes)", repo_path, fc, bs) 372 373 return PackResult(shards=shards, skipped_keys=skipped)
Pack each (scene, phase, modality) triple in scan into one tar.
Reusing ScanResult from the scanner avoids re-walking the
source tree — the scan already has the file counts, byte totals, and
structure we need. The packer just translates each (scene, phase,
modality) entry into a tar at the right repo path.
429def pack_objects_meta(plan: PackPlan, scan: ScanResult) -> List[SharedArtefact]: 430 """Build the ``objects_meta/`` layer. 431 432 For every SOS scene with a ``questionnaire.txt`` at its scene root, 433 parse the file and write a normalised 434 ``objects_meta/<object_id>/questionnaire.json`` under the staging 435 dir. ``object_id`` is exactly the SOS scene directory name (the 436 canonical key MOS scenes use in ``mask_to_object.json``). 437 438 Also write ``objects_meta/_index.json``: a small lookup table 439 listing every known object_id. The downloader uses this to verify 440 that mask_to_object references resolve. 441 """ 442 out: List[SharedArtefact] = [] 443 objects_meta = plan.staging_root / "objects_meta" 444 objects_meta.mkdir(parents=True, exist_ok=True) 445 446 object_ids: list[str] = [] 447 for scene in scan.scenes: 448 if scene.scene_type is not SceneType.SINGLE_OBJECT: 449 continue 450 sos_root = plan.src_root / "sos" / scene.scene_id 451 questionnaire_src = sos_root / "questionnaire.txt" 452 if not questionnaire_src.is_file(): 453 log.warning("SOS scene %s has no questionnaire.txt; skipping", scene.scene_id) 454 continue 455 456 parsed = _parse_questionnaire( 457 questionnaire_src.read_text(encoding="utf-8"), 458 ) 459 payload = { 460 "object_id": scene.scene_id, 461 "source": "sos/" + scene.scene_id + "/questionnaire.txt", 462 "questions": parsed, 463 } 464 repo_path = f"objects_meta/{scene.scene_id}/questionnaire.json" 465 out_path = plan.staging_root / repo_path 466 out_path.parent.mkdir(parents=True, exist_ok=True) 467 if out_path.exists() and not plan.overwrite: 468 raise DatasetError( 469 f"refusing to overwrite existing object meta: {out_path}", 470 hint="Pass overwrite=True or remove the staging dir.", 471 ) 472 rendered = json.dumps(payload, indent=2, sort_keys=True) 473 out_path.write_text(rendered, encoding="utf-8") 474 out.append( 475 SharedArtefact( 476 object_id=scene.scene_id, 477 repo_path=repo_path, 478 total_bytes=out_path.stat().st_size, 479 ) 480 ) 481 object_ids.append(scene.scene_id) 482 log.info("wrote %s (%d questions)", repo_path, len(parsed)) 483 484 index_path = objects_meta / "_index.json" 485 index_path.write_text( 486 json.dumps({"object_ids": sorted(object_ids)}, indent=2, sort_keys=True), 487 encoding="utf-8", 488 ) 489 return out
Build the objects_meta/ layer.
For every SOS scene with a questionnaire.txt at its scene root,
parse the file and write a normalised
objects_meta/<object_id>/questionnaire.json under the staging
dir. object_id is exactly the SOS scene directory name (the
canonical key MOS scenes use in mask_to_object.json).
Also write objects_meta/_index.json: a small lookup table
listing every known object_id. The downloader uses this to verify
that mask_to_object references resolve.
SCHEMA_VERSION =
'v1'
@dataclass(frozen=True)
class
ManifestPaths:
80@dataclass(frozen=True) 81class ManifestPaths: 82 """Where the manifest writer dropped its outputs.""" 83 84 parquet_path: Path 85 current_json_path: Path
Where the manifest writer dropped its outputs.
def
build_frame_manifest( scan: ScanResult, pack: PackResult, out_dir: pathlib.Path, splits: Optional[Mapping[str, str]] = None, label_versions: Optional[Mapping[str, str]] = None) -> ManifestPaths:
236def build_frame_manifest( 237 scan: ScanResult, 238 pack: PackResult, 239 out_dir: Path, 240 splits: Optional[Mapping[str, str]] = None, 241 label_versions: Optional[Mapping[str, str]] = None, 242) -> ManifestPaths: 243 """Write ``manifest/frames_<v>.parquet`` and ``manifest/current.json``. 244 245 Parameters 246 ---------- 247 scan : ScanResult 248 Output of :func:`scan_capture_root`. 249 pack : PackResult 250 Output of :func:`pack_capture_tree`. 251 out_dir : Path 252 Where to write ``frames_<v>.parquet`` and ``current.json``. Both 253 end up under ``out_dir / "manifest" / ...``. 254 splits : Mapping[str, str], optional 255 ``{scene_id: "easy" | "medium" | "hard"}`` for multi-object scenes. 256 Single-object scenes are always assigned ``None``. 257 label_versions : Mapping[str, str], optional 258 Logical-label → version map written verbatim to ``current.json``. 259 Defaults to ``{"masks": "v1", "masks_aux": "v1", "sam2_meta": "v1"}``. 260 """ 261 pa, pq = _arrow() 262 splits = dict(splits or {}) 263 label_versions = dict( 264 label_versions 265 or { 266 MASKS: "v1", 267 MASKS_AUX: "v1", 268 "sam2_meta": "v1", 269 CAM_POSE: "v1", 270 } 271 ) 272 out_dir = Path(out_dir) 273 (out_dir / "manifest").mkdir(parents=True, exist_ok=True) 274 275 shard_index = _index_shards_by_phase(pack) 276 277 # Build the table column-by-column for memory efficiency. 278 cols: Dict[str, list] = { 279 "scene_id": [], 280 "scene_type": [], 281 "phase": [], 282 "frame_idx": [], 283 "frame_filename": [], 284 "split": [], 285 } 286 for has_col, shard_col in _COLUMNS.values(): 287 cols[has_col] = [] 288 cols[shard_col] = [] 289 290 for scene in scan.scenes: 291 for phase in scene.phases: 292 shards_here = shard_index.get( 293 (scene.scene_type, scene.scene_id, phase.phase_index), {} 294 ) 295 filenames = _frame_filenames_for( 296 scan, 297 scene.scene_type, 298 scene.scene_id, 299 phase.phase_index, 300 ) 301 # Ego scenes reuse their MOS sibling's split/difficulty tier — 302 # same scene_id (e.g. "scene20.su.checkerboard"), no separate 303 # ego difficulty scoring. SOS (single-object, no scene-level 304 # difficulty concept) still gets None. 305 split_label = ( 306 splits.get(scene.scene_id) 307 if scene.scene_type in (SceneType.MULTI_OBJECT, SceneType.EGO) 308 else None 309 ) 310 for idx, fname in enumerate(filenames): 311 cols["scene_id"].append(scene.scene_id) 312 cols["scene_type"].append(scene.scene_type.value) 313 cols["phase"].append(phase.phase_index) 314 cols["frame_idx"].append(idx) 315 cols["frame_filename"].append(fname) 316 cols["split"].append(split_label) 317 for modality, (has_col, shard_col) in _COLUMNS.items(): 318 shard = shards_here.get(modality) 319 cols[has_col].append(shard is not None) 320 cols[shard_col].append(shard.repo_path if shard else None) 321 322 schema = pa.schema( 323 [ 324 ("scene_id", pa.string()), 325 ("scene_type", pa.string()), 326 ("phase", pa.int32()), 327 ("frame_idx", pa.int32()), 328 ("frame_filename", pa.string()), 329 ("split", pa.string()), 330 *[(c, pa.bool_()) for (c, _) in _COLUMNS.values()], 331 *[(c, pa.string()) for (_, c) in _COLUMNS.values()], 332 ] 333 ) 334 table = pa.table({k: cols[k] for k in schema.names}, schema=schema) 335 336 parquet_path = out_dir / "manifest" / f"frames_{SCHEMA_VERSION}.parquet" 337 pq.write_table(table, parquet_path, compression="zstd") 338 339 current_path = out_dir / "manifest" / "current.json" 340 modality_extensions = _detect_modality_extensions(scan) 341 342 # Preserve any pre-existing current.json keys we do not own. The live 343 # IRVLUTD/RPX repo's current.json carries additional top-level blocks 344 # (``manifests``, ``sos``, ``mos``, ``metadata_versions``, …) that 345 # the teammate's separate manifest tooling writes. A naive overwrite 346 # would strip them and break downstream consumers; instead we read 347 # the existing payload (if any), update only the keys we are 348 # authoritative for, and write the merged result back. 349 # 350 # The keys we own (and will overwrite on every run): 351 # 352 # * ``label_versions`` — from the ``--label-version`` flag 353 # * ``schema_version`` — fixed at this module's SCHEMA_VERSION 354 # * ``modality_extensions`` — sniffed from the scanned source tree 355 # 356 # Everything else is left untouched. 357 OWNED_KEYS = {"label_versions", "schema_version", "modality_extensions"} 358 if current_path.is_file(): 359 try: 360 current_payload: Dict[str, Any] = json.loads( 361 current_path.read_text(encoding="utf-8"), 362 ) 363 if not isinstance(current_payload, dict): 364 current_payload = {} 365 except json.JSONDecodeError: 366 current_payload = {} 367 else: 368 current_payload = {} 369 current_payload["label_versions"] = label_versions 370 current_payload["schema_version"] = SCHEMA_VERSION 371 if modality_extensions: 372 current_payload["modality_extensions"] = modality_extensions 373 # Don't carry stale modality_extensions if it has nothing to say 374 # AND no prior writer set it. 375 elif "modality_extensions" in current_payload and not current_payload["modality_extensions"]: 376 current_payload.pop("modality_extensions") 377 preserved = sorted(k for k in current_payload if k not in OWNED_KEYS) 378 if preserved: 379 log.info( 380 "current.json: preserved %d unmanaged top-level key(s): %s", 381 len(preserved), 382 ", ".join(preserved), 383 ) 384 current_path.write_text( 385 json.dumps(current_payload, indent=2), 386 encoding="utf-8", 387 ) 388 389 # Persist per-tar SHA-256 to manifest/checksums.json. The packer 390 # computes these as part of pack_capture_tree; the downstream 391 # download-side verification in hub._extract_snapshot_tars reads 392 # this file and aborts on any tar whose bytes don't match. Tar 393 # entries whose sha256 is None (rare — only when re-running 394 # manifest without a fresh pack) are still written with a null 395 # value so the absence is explicit, not silent. 396 checksums_path = out_dir / "manifest" / "checksums.json" 397 checksums = {s.repo_path: s.sha256 for s in pack.shards} 398 checksums_path.write_text( 399 json.dumps( 400 {"sha256": checksums, "algorithm": "sha256"}, 401 indent=2, 402 sort_keys=True, 403 ), 404 encoding="utf-8", 405 ) 406 407 log.info("wrote manifest: %d rows → %s", len(table), parquet_path) 408 n_with_sha = sum(1 for v in checksums.values() if v) 409 log.info( 410 "wrote checksums.json: %d/%d tar shards have SHA-256 → %s", 411 n_with_sha, 412 len(checksums), 413 checksums_path, 414 ) 415 416 # Per-file SHA-256: walk every packed tar, stream each member through 417 # SHA-256, and persist the mapping at manifest/file_checksums.json. 418 # This is the foundation for ``verify_dataset(snapshot_root)`` to 419 # detect local-disk corruption between extract and load — the case 420 # the tar-level SHA-256 cannot guard against. 421 file_checksums = _compute_per_file_checksums(out_dir, pack) 422 file_checksums_path = out_dir / "manifest" / "file_checksums.json" 423 file_checksums_path.write_text( 424 json.dumps( 425 {"algorithm": "sha256", "sha256": file_checksums}, 426 indent=2, 427 sort_keys=True, 428 ), 429 encoding="utf-8", 430 ) 431 log.info( 432 "wrote file_checksums.json: %d extracted-file SHA-256 entries → %s", 433 len(file_checksums), 434 file_checksums_path, 435 ) 436 437 if modality_extensions: 438 log.info( 439 "current.json modality_extensions: %s", 440 ", ".join(f"{k}={v}" for k, v in sorted(modality_extensions.items())), 441 ) 442 return ManifestPaths(parquet_path=parquet_path, current_json_path=current_path)
Write manifest/frames_<v>.parquet and manifest/current.json.
Parameters
scan : ScanResult
Output of scan_capture_root().
pack : PackResult
Output of pack_capture_tree().
out_dir : Path
Where to write frames_<v>.parquet and current.json. Both
end up under out_dir / "manifest" / ....
splits : Mapping[str, str], optional
{scene_id: "easy" | "medium" | "hard"} for multi-object scenes.
Single-object scenes are always assigned None.
label_versions : Mapping[str, str], optional
Logical-label → version map written verbatim to current.json.
Defaults to {"masks": "v1", "masks_aux": "v1", "sam2_meta": "v1"}.
def
read_frame_manifest(parquet_path: pathlib.Path):
501def read_frame_manifest(parquet_path: Path): 502 """Read a manifest Parquet back as a pyarrow Table (test/debug aid).""" 503 _, pq = _arrow() 504 return pq.read_table(parquet_path)
Read a manifest Parquet back as a pyarrow Table (test/debug aid).
def
write_split_manifests( staging_root: str | pathlib.Path, *, tasks: Optional[Iterable[str]] = None, splits: Iterable[str] = ('easy', 'medium', 'hard')) -> Dict[Tuple[str, str], pathlib.Path]:
597def write_split_manifests( 598 staging_root: str | Path, 599 *, 600 tasks: Optional[Iterable[str]] = None, 601 splits: Iterable[str] = _SPLITS, 602) -> Dict[Tuple[str, str], Path]: 603 """Write per-(task, split) manifest JSONs under 604 ``<staging_root>/manifests/<task>/<split>.json``. 605 606 Returns ``{(recipe_key, split): path}`` for every JSON written. 607 Empty (task, split) combinations (no rows after filtering, or specs 608 that intentionally emit nothing — e.g. VQA) are skipped silently. 609 """ 610 try: 611 import pandas as pd # noqa: PLC0415 612 except ImportError as e: 613 raise ImportError( 614 "pandas is required for split_manifests; install via `pip install pandas`." 615 ) from e 616 617 staging_root = Path(staging_root) 618 parquet = _parquet_path(staging_root) 619 df = pd.read_parquet(parquet) 620 621 # Authoritative scene-wise splits. 622 scene_split = _scene_wise_splits(df) 623 df = df.drop(columns=["split"]).merge( 624 scene_split.rename("split"), 625 left_on="scene_id", 626 right_index=True, 627 ) 628 629 # Read per-modality extensions from current.json (v2+ datasets). 630 # Datasets without this field (v1) fall through to the hardcoded 631 # defaults in _MODALITY_LAYOUT. 632 extensions = _load_modality_extensions(staging_root) 633 layout_override = _layout_with_overrides(extensions) if extensions else None 634 if extensions: 635 log.info( 636 "split_manifests: using modality_extensions from current.json: %s", 637 ", ".join(f"{k}={v}" for k, v in sorted(extensions.items())), 638 ) 639 640 chosen = list(tasks) if tasks is not None else list(_TASK_SPECS.keys()) 641 manifests_dir = staging_root / "manifests" 642 written: Dict[Tuple[str, str], Path] = {} 643 644 global _active_layout 645 _active_layout = layout_override 646 try: 647 written = _write_loop( 648 chosen=chosen, 649 df=df, 650 splits=splits, 651 manifests_dir=manifests_dir, 652 ) 653 finally: 654 _active_layout = None 655 656 return written
Write per-(task, split) manifest JSONs under
<staging_root>/manifests/<task>/<split>.json.
Returns {(recipe_key, split): path} for every JSON written.
Empty (task, split) combinations (no rows after filtering, or specs
that intentionally emit nothing — e.g. VQA) are skipped silently.
DEFAULT_IGNORE_PATTERNS =
('.DS_Store', '__pycache__', '*.pyc', '*.tmp', '*.part', '.huggingface/*')
@dataclass(frozen=True)
class
UploadPlan:
41@dataclass(frozen=True) 42class UploadPlan: 43 """Inputs needed to push a staging directory to a HF dataset repo.""" 44 45 staging_root: Path 46 repo_id: str = DEFAULT_REPO_ID 47 revision: str = "main" 48 private: bool = True 49 commit_message: str = "RPX dataset hub upload" 50 create_if_missing: bool = True 51 use_large_folder: bool = True 52 allow_patterns: Optional[Sequence[str]] = None 53 ignore_patterns: Sequence[str] = DEFAULT_IGNORE_PATTERNS 54 dry_run: bool = False
Inputs needed to push a staging directory to a HF dataset repo.
UploadPlan( staging_root: pathlib.Path, repo_id: str = 'IRVLUTD/RPX', revision: str = 'main', private: bool = True, commit_message: str = 'RPX dataset hub upload', create_if_missing: bool = True, use_large_folder: bool = True, allow_patterns: Optional[Sequence[str]] = None, ignore_patterns: Sequence[str] = ('.DS_Store', '__pycache__', '*.pyc', '*.tmp', '*.part', '.huggingface/*'), dry_run: bool = False)
@dataclass(frozen=True)
class
UploadResult:
57@dataclass(frozen=True) 58class UploadResult: 59 """Reported back so callers can log / verify.""" 60 61 repo_id: str 62 revision: str 63 files_planned: int 64 bytes_planned: int 65 commit_url: Optional[str] = None 66 skipped: List[Path] = field(default_factory=list)
Reported back so callers can log / verify.
def
upload_paths( paths: Sequence[pathlib.Path], staging_root: pathlib.Path, repo_id: str = 'IRVLUTD/RPX', revision: str = 'main', private: bool = True, commit_message: str = 'RPX dataset hub: partial upload', dry_run: bool = False) -> UploadResult:
213def upload_paths( 214 paths: Sequence[Path], 215 staging_root: Path, 216 repo_id: str = DEFAULT_REPO_ID, 217 revision: str = "main", 218 private: bool = True, 219 commit_message: str = "RPX dataset hub: partial upload", 220 dry_run: bool = False, 221) -> UploadResult: 222 """Upload a specific list of files (relative to ``staging_root``). 223 224 Useful for shipping a single modality or just the manifest after a 225 partial re-pack. Internally builds an ``allow_patterns`` list and 226 delegates to :func:`upload_staging` (which still calls 227 ``upload_large_folder``/``upload_folder`` under the hood). 228 """ 229 rels = [] 230 total = 0 231 for p in paths: 232 rel = p.relative_to(staging_root).as_posix() 233 rels.append(rel) 234 total += p.stat().st_size 235 236 plan = UploadPlan( 237 staging_root=staging_root, 238 repo_id=repo_id, 239 revision=revision, 240 private=private, 241 commit_message=commit_message, 242 allow_patterns=tuple(rels), 243 use_large_folder=False, 244 dry_run=dry_run, 245 ) 246 return upload_staging(plan)
Upload a specific list of files (relative to staging_root).
Useful for shipping a single modality or just the manifest after a
partial re-pack. Internally builds an allow_patterns list and
delegates to upload_staging() (which still calls
upload_large_folder/upload_folder under the hood).
111def upload_staging(plan: UploadPlan) -> UploadResult: 112 """Push everything under ``plan.staging_root`` to the HF dataset repo. 113 114 Behaviour: 115 116 * If the repo doesn't exist and ``create_if_missing`` is True, it is 117 created (private by default, dataset type). 118 * If ``dry_run`` is True, no network call is made — we walk the tree 119 and report what *would* be uploaded. 120 * Otherwise, ``upload_large_folder`` is preferred for true resumable 121 multi-commit pushes; ``upload_folder`` is the fallback for tiny 122 mock uploads where the multi-commit overhead isn't worth it. 123 """ 124 if not plan.staging_root.is_dir(): 125 raise ConfigError( 126 f"staging dir does not exist: {plan.staging_root}", 127 hint="Run pack_capture_tree() and build_frame_manifest() first.", 128 ) 129 130 files, total = _enumerate_uploadable( 131 plan.staging_root, 132 plan.ignore_patterns, 133 plan.allow_patterns, 134 ) 135 log.info( 136 "upload plan: %d files, %d bytes → %s (revision=%s, dry_run=%s)", 137 len(files), 138 total, 139 plan.repo_id, 140 plan.revision, 141 plan.dry_run, 142 ) 143 144 if plan.dry_run: 145 return UploadResult( 146 repo_id=plan.repo_id, 147 revision=plan.revision, 148 files_planned=len(files), 149 bytes_planned=total, 150 ) 151 152 hf = _hub() 153 api = hf.HfApi() 154 155 if plan.create_if_missing: 156 try: 157 api.create_repo( 158 repo_id=plan.repo_id, 159 repo_type="dataset", 160 private=plan.private, 161 exist_ok=True, 162 ) 163 except Exception as e: 164 raise DownloadError( 165 f"failed to create or verify dataset repo {plan.repo_id}: {e}", 166 hint="Run `hf auth login` and confirm you can push to the org.", 167 ) from e 168 169 common_kwargs = dict( 170 repo_id=plan.repo_id, 171 repo_type="dataset", 172 folder_path=str(plan.staging_root), 173 revision=plan.revision, 174 ignore_patterns=list(plan.ignore_patterns), 175 ) 176 if plan.allow_patterns: 177 common_kwargs["allow_patterns"] = list(plan.allow_patterns) 178 179 try: 180 if plan.use_large_folder and hasattr(api, "upload_large_folder"): 181 # `upload_large_folder` does its own multi-commit chunking 182 # and writes checkpoint state under .huggingface/ so re-runs 183 # resume cleanly. 184 api.upload_large_folder(**common_kwargs) 185 commit_url = None 186 else: 187 commit = api.upload_folder( 188 commit_message=plan.commit_message, 189 **common_kwargs, 190 ) 191 commit_url = getattr(commit, "commit_url", None) 192 except Exception as e: 193 raise DownloadError( 194 f"upload to {plan.repo_id} failed: {e}", 195 hint=( 196 "Re-run to resume; HF stores checkpoint state under " 197 "<staging>/.huggingface/. If the failure is auth, run " 198 "`hf auth login` and confirm push access to the org." 199 ), 200 ) from e 201 202 return UploadResult( 203 repo_id=plan.repo_id, 204 revision=plan.revision, 205 files_planned=len(files), 206 bytes_planned=total, 207 commit_url=commit_url, 208 )
Push everything under plan.staging_root to the HF dataset repo.
Behaviour:
- If the repo doesn't exist and
create_if_missingis True, it is created (private by default, dataset type). - If
dry_runis True, no network call is made — we walk the tree and report what would be uploaded. - Otherwise,
upload_large_folderis preferred for true resumable multi-commit pushes;upload_folderis the fallback for tiny mock uploads where the multi-commit overhead isn't worth it.
@dataclass(frozen=True)
class
CroissantPatch:
46@dataclass(frozen=True) 47class CroissantPatch: 48 """Knobs for the staging step.""" 49 50 repo_id: str = DEFAULT_REPO_ID 51 version: str = "1.0.0" 52 cite_as: Optional[str] = None 53 annotate_layout: bool = True
Knobs for the staging step.
def
stage_croissant( staging_root: pathlib.Path, src: Optional[pathlib.Path] = None, patch: Optional[CroissantPatch] = None, overwrite: bool = False) -> pathlib.Path:
70def stage_croissant( 71 staging_root: Path, 72 src: Optional[Path] = None, 73 patch: Optional[CroissantPatch] = None, 74 overwrite: bool = False, 75) -> Path: 76 """Copy + patch the Croissant JSON into ``<staging>/rpx_croissant.json``. 77 78 Returns the destination path. 79 """ 80 src = Path(src) if src else _default_croissant_src() 81 patch = patch or CroissantPatch() 82 if not src.is_file(): 83 raise ConfigError( 84 f"croissant source JSON does not exist: {src}", 85 hint=( 86 "Pass src= explicitly, or check that " 87 "croissant_template.json is present in the package " 88 "in the repo." 89 ), 90 ) 91 92 out_path = Path(staging_root) / "rpx_croissant.json" 93 if out_path.exists() and not overwrite: 94 raise DatasetError( 95 f"refusing to overwrite existing croissant: {out_path}", 96 hint="Pass overwrite=True or remove the file first.", 97 ) 98 99 payload = json.loads(src.read_text(encoding="utf-8")) 100 101 payload["url"] = _hf_url(patch.repo_id) 102 payload["version"] = patch.version 103 if patch.cite_as is not None: 104 payload["citeAs"] = patch.cite_as 105 106 if patch.annotate_layout: 107 desc = payload.get("description", "") 108 if _LAYOUT_NOTE.strip() not in desc: 109 payload["description"] = desc + _LAYOUT_NOTE 110 111 out_path.parent.mkdir(parents=True, exist_ok=True) 112 out_path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8") 113 log.info("staged croissant → %s (%d bytes)", out_path, out_path.stat().st_size) 114 return out_path
Copy + patch the Croissant JSON into <staging>/rpx_croissant.json.
Returns the destination path.
@dataclass(frozen=True)
class
CardSpec:
58@dataclass(frozen=True) 59class CardSpec: 60 """Knobs for the card generator.""" 61 62 repo_id: str = DEFAULT_REPO_ID 63 pretty_name: str = DEFAULT_PRETTY_NAME 64 license: str = DEFAULT_LICENSE 65 task_categories: tuple[str, ...] = DEFAULT_TASK_CATEGORIES 66 tags: tuple[str, ...] = DEFAULT_TAGS 67 paper_url: Optional[str] = None 68 code_url: Optional[str] = "https://github.com/IRVLUTD/RPX"
Knobs for the card generator.
CardSpec( repo_id: str = 'IRVLUTD/RPX', pretty_name: str = 'RPX: Robot Perception X', license: str = 'cc-by-4.0', task_categories: tuple[str, ...] = ('image-segmentation', 'depth-estimation', 'object-detection', 'visual-question-answering'), tags: tuple[str, ...] = ('robotics', 'embodied-ai', 'rgb-d', 'benchmark', 'perception', 'manipulation', 'stereo'), paper_url: Optional[str] = None, code_url: Optional[str] = 'https://github.com/IRVLUTD/RPX')
def
write_dataset_card( staging_root: pathlib.Path, scan: ScanResult, spec: Optional[CardSpec] = None, splits: Optional[Mapping[str, str]] = None, label_versions: Optional[Mapping[str, str]] = None, overwrite: bool = False) -> pathlib.Path:
313def write_dataset_card( 314 staging_root: Path, 315 scan: ScanResult, 316 spec: Optional[CardSpec] = None, 317 splits: Optional[Mapping[str, str]] = None, 318 label_versions: Optional[Mapping[str, str]] = None, 319 overwrite: bool = False, 320) -> Path: 321 """Render and write ``<staging_root>/README.md``. 322 323 Returns the destination path. 324 """ 325 spec = spec or CardSpec() 326 out_path = Path(staging_root) / "README.md" 327 if out_path.exists() and not overwrite: 328 raise DatasetError( 329 f"refusing to overwrite existing dataset card: {out_path}", 330 hint="Pass overwrite=True or remove the file first.", 331 ) 332 333 contents = ( 334 _frontmatter(spec, scan, splits) 335 + "\n" 336 + _body( 337 spec, 338 scan, 339 splits, 340 label_versions, 341 ) 342 ) 343 out_path.parent.mkdir(parents=True, exist_ok=True) 344 out_path.write_text(contents, encoding="utf-8") 345 log.info("wrote dataset card → %s (%d bytes)", out_path, out_path.stat().st_size) 346 return out_path
Render and write <staging_root>/README.md.
Returns the destination path.
@dataclass(frozen=True)
class
DownloadResult:
70@dataclass(frozen=True) 71class DownloadResult: 72 """Outcome of one selective pull.""" 73 74 repo_id: str 75 revision: Optional[str] 76 task: str 77 split: Optional[str] 78 scene_type: SceneType 79 local_dir: Path 80 allow_patterns: List[str] 81 matched_scenes: List[str] 82 bytes_fetched: int 83 files_fetched: int 84 cache_hits: int = 0 85 manifest_table: Optional["pa.Table"] = None
Outcome of one selective pull.
DownloadResult( repo_id: str, revision: Optional[str], task: str, split: Optional[str], scene_type: SceneType, local_dir: pathlib.Path, allow_patterns: List[str], matched_scenes: List[str], bytes_fetched: int, files_fetched: int, cache_hits: int = 0, manifest_table: Optional[pyarrow.lib.Table] = None)
scene_type: SceneType
def
download_for_task( task: str, split: Optional[str] = None, repo_id: str = 'IRVLUTD/RPX', *, revision: Optional[str] = None, cache_dir: Optional[pathlib.Path] = None, extra_modalities: Sequence[str] = (), label_versions: Optional[Mapping[str, str]] = None, scene_type: Optional[SceneType] = None, max_workers: int = 8, return_manifest: bool = True) -> DownloadResult:
235def download_for_task( 236 task: str, 237 split: Optional[str] = None, 238 repo_id: str = DEFAULT_REPO_ID, 239 *, 240 revision: Optional[str] = None, 241 cache_dir: Optional[Path] = None, 242 extra_modalities: Sequence[str] = (), 243 label_versions: Optional[Mapping[str, str]] = None, 244 scene_type: Optional[SceneType] = None, 245 max_workers: int = 8, 246 return_manifest: bool = True, 247) -> DownloadResult: 248 """Pull just the tar shards needed to run ``task`` on ``split``. 249 250 Parameters 251 ---------- 252 task : str 253 Recipe key from :data:`MULTI_OBJECT_TASK_RECIPES` or 254 :data:`SINGLE_OBJECT_TASK_RECIPES`. 255 split : str, optional 256 ``"easy" | "medium" | "hard"``. Required for multi-object tasks, 257 forbidden for single-object tasks (which have no splits). 258 repo_id : str 259 HuggingFace dataset repo (default: ``IRVLUTD/RPX``). 260 extra_modalities : Sequence[str] 261 Add modalities beyond the recipe's defaults (e.g. ``["masks_aux"]`` 262 if you want the visualisation overlays alongside ``masks``). 263 label_versions : Mapping[str, str], optional 264 Override which label version to pull for one or more modalities. 265 Default behaviour is to follow the repo's ``manifest/current.json``. 266 scene_type : SceneType, optional 267 Disambiguates when a task name exists in both recipe tables. 268 269 Returns 270 ------- 271 DownloadResult 272 ``local_dir`` is the local snapshot root; combine with 273 ``allow_patterns`` and ``manifest_table`` to feed a downstream 274 loader. ``bytes_fetched`` reports the size on disk after the pull 275 (for true delta accounting prefer your own ``du -sh``). 276 """ 277 hf = _hub() 278 recipe: TaskRecipe = resolve_recipe(task, scene_type=scene_type) 279 280 if recipe.scene_type is SceneType.SINGLE_OBJECT and split is not None: 281 raise ConfigError( 282 f"task {task!r} targets single-object scenes, which have no splits.", 283 hint="Drop the --split argument for single-object recipes.", 284 ) 285 if recipe.scene_type in _SPLIT_BEARING_TYPES and split is None: 286 raise ConfigError( 287 f"task {task!r} targets {recipe.scene_type.value} scenes — pick a split.", 288 hint="Pass split='easy' | 'medium' | 'hard'.", 289 ) 290 291 parquet_path, current_path = _fetch_manifest( 292 hf, 293 repo_id, 294 revision, 295 cache_dir, 296 ) 297 current = json.loads(current_path.read_text(encoding="utf-8")) 298 _, pq = _arrow() 299 table = pq.read_table(parquet_path) 300 sliced = _filter_manifest(table, recipe.scene_type, split) 301 302 matched_scenes: Set[str] = set(sliced["scene_id"].to_pylist()) 303 if not matched_scenes: 304 raise DownloadError( 305 f"No scenes matched task={task!r} split={split!r} on {repo_id}.", 306 hint=( 307 "Verify the manifest has rows for this scene_type/split. " 308 "If the dataset was uploaded without a `splits=...` argument " 309 "to build_frame_manifest, multi-object rows will all have " 310 "split=null." 311 ), 312 ) 313 314 modalities = set(recipe.all_modalities()) | set(extra_modalities) 315 explicit = dict(label_versions or {}) 316 patterns = _build_allow_patterns( 317 matched_scenes, 318 recipe.scene_type, 319 modalities, 320 current, 321 explicit, 322 ) 323 patterns.extend(_shared_artefact_patterns(modalities)) 324 325 log.info( 326 "download plan: task=%s split=%s scenes=%d modalities=%s patterns=%d", 327 task, 328 split, 329 len(matched_scenes), 330 sorted(modalities), 331 len(patterns), 332 ) 333 334 try: 335 local_dir = hf.snapshot_download( 336 repo_id=repo_id, 337 repo_type="dataset", 338 allow_patterns=patterns, 339 revision=revision, 340 cache_dir=str(cache_dir) if cache_dir else None, 341 max_workers=max_workers, 342 ) 343 except Exception as e: 344 raise DownloadError( 345 f"snapshot_download failed for {repo_id}: {e}", 346 hint=( 347 "Re-run to retry; HF caches partial downloads. If auth is the " 348 "issue, run `hf auth login` and confirm read access." 349 ), 350 ) from e 351 352 local = Path(local_dir) 353 354 # Mirror the post-fetch state that hub.download_split produces: extract 355 # every tar shard to flat per-modality dirs so downstream consumers 356 # (loader, viz, custom scripts) hit pure disk reads instead of opening 357 # and seeking tars per frame. Idempotent — files already on disk skip. 358 from ..hub import _extract_snapshot_tars 359 _extract_snapshot_tars(local) 360 361 files = [p for p in local.rglob("*") if p.is_file()] 362 bytes_ = sum(p.stat().st_size for p in files) 363 364 return DownloadResult( 365 repo_id=repo_id, 366 revision=revision, 367 task=task, 368 split=split, 369 scene_type=recipe.scene_type, 370 local_dir=local, 371 allow_patterns=patterns, 372 matched_scenes=sorted(matched_scenes), 373 bytes_fetched=bytes_, 374 files_fetched=len(files), 375 cache_hits=0, # HF doesn't expose per-call hit count 376 manifest_table=sliced if return_manifest else None, 377 )
Pull just the tar shards needed to run task on split.
Parameters
task : str
Recipe key from MULTI_OBJECT_TASK_RECIPES or
SINGLE_OBJECT_TASK_RECIPES.
split : str, optional
"easy" | "medium" | "hard". Required for multi-object tasks,
forbidden for single-object tasks (which have no splits).
repo_id : str
HuggingFace dataset repo (default: IRVLUTD/RPX).
extra_modalities : Sequence[str]
Add modalities beyond the recipe's defaults (e.g. ["masks_aux"]
if you want the visualisation overlays alongside masks).
label_versions : Mapping[str, str], optional
Override which label version to pull for one or more modalities.
Default behaviour is to follow the repo's manifest/current.json.
scene_type : SceneType, optional
Disambiguates when a task name exists in both recipe tables.
Returns
DownloadResult
local_dir is the local snapshot root; combine with
allow_patterns and manifest_table to feed a downstream
loader. bytes_fetched reports the size on disk after the pull
(for true delta accounting prefer your own du -sh).
SPLIT_TIER_FILES =
('easy.txt', 'medium.txt', 'hard.txt')
SCENE_SPLITS_JSON =
'scene_splits.json'
@dataclass(frozen=True)
class
StagedFile:
51@dataclass(frozen=True) 52class StagedFile: 53 """One file the staging step copied / wrote.""" 54 55 repo_path: str # path relative to staging root, forward-slashes 56 src: Optional[Path] 57 bytes_: int
One file the staging step copied / wrote.
def
load_scene_splits(staging_root: pathlib.Path) -> dict[str, str]:
181def load_scene_splits(staging_root: Path) -> dict[str, str]: 182 """Read ``<staging_root>/splits/scene_splits.json`` and return a 183 flat ``{scene_id: tier}`` mapping. Convenience for downstream steps 184 (manifest builder, dataset card) that need to look up a scene's tier. 185 186 Accepts both the canonical schema with a top-level ``"splits"`` key 187 holding ``{tier: [scene_id, ...]}`` and a flat ``{scene_id: tier}`` 188 shape. 189 """ 190 path = Path(staging_root) / "splits" / SCENE_SPLITS_JSON 191 if not path.is_file(): 192 raise ConfigError( 193 f"scene_splits.json not found at {path}", 194 hint="Run stage_splits() before load_scene_splits().", 195 ) 196 payload = json.loads(path.read_text(encoding="utf-8")) 197 if "splits" in payload and isinstance(payload["splits"], dict): 198 return {sid: tier for tier, ids in payload["splits"].items() for sid in ids} 199 return dict(payload)
Read <staging_root>/splits/scene_splits.json and return a
flat {scene_id: tier} mapping. Convenience for downstream steps
(manifest builder, dataset card) that need to look up a scene's tier.
Accepts both the canonical schema with a top-level "splits" key
holding {tier: [scene_id, ...]} and a flat {scene_id: tier}
shape.
def
stage_paths( paths: Sequence[pathlib.Path], staging_root: pathlib.Path, repo_subdir: str = '', overwrite: bool = False) -> List[StagedFile]:
138def stage_paths( 139 paths: Sequence[Path], 140 staging_root: Path, 141 repo_subdir: str = "", 142 overwrite: bool = False, 143) -> List[StagedFile]: 144 """Generic helper: copy a list of files into ``<staging>/<repo_subdir>/``. 145 146 Useful for ad-hoc staging of small files (Croissant JSON, custom 147 docs) without writing a dedicated module. ``repo_subdir`` may be 148 empty to land files at the staging root. 149 """ 150 staging_root = Path(staging_root) 151 target = staging_root / repo_subdir if repo_subdir else staging_root 152 target.mkdir(parents=True, exist_ok=True) 153 154 staged: List[StagedFile] = [] 155 for src in paths: 156 src = Path(src) 157 if not src.is_file(): 158 raise ConfigError( 159 f"source file does not exist: {src}", 160 hint="Check the path; staging does not create files.", 161 ) 162 dst = target / src.name 163 if dst.exists() and not overwrite: 164 raise DatasetError( 165 f"refusing to overwrite existing staged file: {dst}", 166 hint="Pass overwrite=True or remove the file first.", 167 ) 168 shutil.copy2(src, dst) 169 rel = (Path(repo_subdir) / src.name).as_posix() if repo_subdir else src.name 170 staged.append( 171 StagedFile( 172 repo_path=rel, 173 src=src, 174 bytes_=dst.stat().st_size, 175 ) 176 ) 177 log.info("staged %s (%d bytes)", rel, dst.stat().st_size) 178 return staged
Generic helper: copy a list of files into <staging>/<repo_subdir>/.
Useful for ad-hoc staging of small files (Croissant JSON, custom
docs) without writing a dedicated module. repo_subdir may be
empty to land files at the staging root.
def
stage_splits( staging_root: pathlib.Path, splits_src: Optional[pathlib.Path] = None, overwrite: bool = False, require_all: bool = True) -> List[StagedFile]:
60def stage_splits( 61 staging_root: Path, 62 splits_src: Optional[Path] = None, 63 overwrite: bool = False, 64 require_all: bool = True, 65) -> List[StagedFile]: 66 """Copy ``splits/{scene_splits.json, easy.txt, medium.txt, hard.txt}`` 67 into ``<staging_root>/splits/``. 68 69 Parameters 70 ---------- 71 staging_root : Path 72 Where the packed/upload-ready tree lives. 73 splits_src : Path, optional 74 Source dir holding the splits files. Defaults to 75 ``benchmark/data/splits/`` next to this package. 76 overwrite : bool 77 If False (default), raises when a destination file already exists. 78 require_all : bool 79 If True (default), raises when any of the four expected files is 80 missing from ``splits_src``. Set False for partial-stage scenarios 81 (e.g., shipping with only the JSON manifest). 82 83 Returns 84 ------- 85 List[StagedFile] 86 One entry per copied file. 87 """ 88 src_root = Path(splits_src) if splits_src else _default_splits_src() 89 if not src_root.is_dir(): 90 raise ConfigError( 91 f"splits source dir does not exist: {src_root}", 92 hint="Splits files are bundled on the HuggingFace dataset " 93 "(splits/{scene_splits.json,easy.txt,medium.txt,hard.txt}); " 94 "make sure you've downloaded them, or pass splits_src= " 95 "pointing at a local copy.", 96 ) 97 98 staging_root = Path(staging_root) 99 out_dir = staging_root / "splits" 100 out_dir.mkdir(parents=True, exist_ok=True) 101 102 expected = (SCENE_SPLITS_JSON, *SPLIT_TIER_FILES) 103 missing = [name for name in expected if not (src_root / name).is_file()] 104 if missing and require_all: 105 raise DatasetError( 106 f"missing splits files in {src_root}: {missing}", 107 hint=( 108 "Re-download the HuggingFace dataset to recover the " 109 "splits/ tree, or pass require_all=False to ship a " 110 "partial splits dir." 111 ), 112 ) 113 114 staged: List[StagedFile] = [] 115 for name in expected: 116 src = src_root / name 117 if not src.is_file(): 118 log.warning("splits source missing: %s (skipping)", src) 119 continue 120 dst = out_dir / name 121 if dst.exists() and not overwrite: 122 raise DatasetError( 123 f"refusing to overwrite existing staged file: {dst}", 124 hint="Pass overwrite=True or remove the file first.", 125 ) 126 shutil.copy2(src, dst) 127 staged.append( 128 StagedFile( 129 repo_path=f"splits/{name}", 130 src=src, 131 bytes_=dst.stat().st_size, 132 ) 133 ) 134 log.info("staged splits/%s (%d bytes)", name, dst.stat().st_size) 135 return staged
Copy splits/{scene_splits.json, easy.txt, medium.txt, hard.txt}
into <staging_root>/splits/.
Parameters
staging_root : Path
Where the packed/upload-ready tree lives.
splits_src : Path, optional
Source dir holding the splits files. Defaults to
benchmark/data/splits/ next to this package.
overwrite : bool
If False (default), raises when a destination file already exists.
require_all : bool
If True (default), raises when any of the four expected files is
missing from splits_src. Set False for partial-stage scenarios
(e.g., shipping with only the JSON manifest).
Returns
List[StagedFile] One entry per copied file.