"""Resume-safe download and preparation of the planned Hypersim expansion.""" from __future__ import annotations import argparse import json import shutil import time from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import asdict from pathlib import Path from mcss.data.download import download_hypersim_members from mcss.data.hypersim import prepare_hypersim from mcss.data.hypersim_expansion import ( PROTOCOL_PARTITIONS, ExpansionScene, build_expansion_plan, select_materialization_scenes, ) from mcss.data.manifest import load_manifest GIB = 1024**3 def _parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--stage", choices=("download", "prepare", "all"), required=True) parser.add_argument( "--metadata", type=Path, default=Path("data/hypersim_raw/official_metadata/metadata_images_split_scene_v1.csv"), ) parser.add_argument( "--observed-subset", type=Path, default=Path("configs/hypersim_subset_v1.csv") ) parser.add_argument("--raw-root", type=Path, default=Path("data/hypersim_raw")) parser.add_argument( "--prepared-root", type=Path, default=Path("data/hypersim_er_prepared") ) parser.add_argument( "--partitions", nargs="+", choices=PROTOCOL_PARTITIONS, default=["train", "val", "diagnostic_test"], ) parser.add_argument("--allow-final-holdout", action="store_true") parser.add_argument("--image-height", type=int, default=128) parser.add_argument("--image-width", type=int, default=160) parser.add_argument("--workers", type=int, default=2) parser.add_argument("--limit", type=int) parser.add_argument("--reserve-gib", type=float, default=100.0) parser.add_argument( "--log", type=Path, default=Path("outputs/hypersim_er_materialization/materialize.jsonl") ) parser.add_argument("--dry-run", action="store_true") return parser def main() -> int: args = _parser().parse_args() if args.workers < 1: raise ValueError("--workers must be positive") if args.limit is not None and args.limit < 1: raise ValueError("--limit must be positive") if args.reserve_gib < 0: raise ValueError("--reserve-gib must be non-negative") if args.image_height < 1 or args.image_width < 1: raise ValueError("image dimensions must be positive") plan = build_expansion_plan(args.metadata, args.observed_subset) scenes = select_materialization_scenes( plan, args.partitions, allow_final_holdout=args.allow_final_holdout, ) if args.limit is not None: scenes = scenes[: args.limit] summary = { "stage": args.stage, "dry_run": args.dry_run, "scene_count": len(scenes), "partitions": list(dict.fromkeys(scene.protocol_partition for scene in scenes)), "metadata_sha256": plan.metadata_sha256, "reserve_bytes": int(args.reserve_gib * GIB), } print(json.dumps({"event": "plan", **summary}, sort_keys=True), flush=True) if args.dry_run: for scene in scenes: print(json.dumps({"event": "scene", **asdict(scene)}, sort_keys=True)) return 0 args.log.parent.mkdir(parents=True, exist_ok=True) failures: list[dict[str, object]] = [] completed = 0 with ThreadPoolExecutor(max_workers=args.workers) as pool: futures = { pool.submit( _materialize_one, scene, stage=args.stage, raw_root=args.raw_root, prepared_root=args.prepared_root, image_size=(args.image_height, args.image_width), reserve_bytes=int(args.reserve_gib * GIB), ): scene for scene in scenes } for future in as_completed(futures): scene = futures[future] try: record = future.result() completed += 1 except Exception as error: # Continue so a transient scene failure is resumable. record = { "event": "failed", "scene_name": scene.scene_name, "protocol_partition": scene.protocol_partition, "error_type": type(error).__name__, "error": str(error), } failures.append(record) _append_jsonl(args.log, record) print( json.dumps( {**record, "progress": f"{completed + len(failures)}/{len(scenes)}"}, sort_keys=True, ), flush=True, ) final = { "event": "summary", **summary, "completed": completed, "failed": len(failures), "free_bytes": shutil.disk_usage(args.raw_root.resolve().anchor).free, } _append_jsonl(args.log, final) print(json.dumps(final, sort_keys=True), flush=True) return 0 if not failures else 1 def _materialize_one( scene: ExpansionScene, *, stage: str, raw_root: Path, prepared_root: Path, image_size: tuple[int, int], reserve_bytes: int, ) -> dict[str, object]: started = time.perf_counter() downloaded = False prepared = False _require_space(raw_root, reserve_bytes) if stage in {"download", "all"} and not _raw_scene_complete(raw_root, scene): download_hypersim_members( scene.scene_name, raw_root, camera=scene.selected_camera, frame_ids=None, ) downloaded = True if stage in {"prepare", "all"}: if not _raw_scene_complete(raw_root, scene): raise FileNotFoundError( f"raw scene is incomplete: {scene.scene_name}/{scene.selected_camera}" ) destination = prepared_root / scene.protocol_partition / scene.scene_name if not _prepared_scene_complete(destination, scene.frame_count): _require_space(prepared_root, reserve_bytes) paths = prepare_hypersim( raw_root, prepared_root / scene.protocol_partition, scenes=[scene.scene_name], camera=scene.selected_camera, image_size=image_size, frame_stride=1, ) if len(paths) != 1: raise RuntimeError(f"expected one prepared manifest for {scene.scene_name}") prepared = True _require_space(raw_root, reserve_bytes) return { "event": "completed", "scene_name": scene.scene_name, "protocol_partition": scene.protocol_partition, "selected_camera": scene.selected_camera, "frame_count": scene.frame_count, "downloaded": downloaded, "prepared": prepared, "elapsed_seconds": time.perf_counter() - started, } def _raw_scene_complete(root: Path, scene: ExpansionScene) -> bool: scene_root = root / scene.scene_name detail = scene_root / "_detail" camera_detail = detail / scene.selected_camera preview = scene_root / "images" / f"scene_{scene.selected_camera}_final_preview" geometry = scene_root / "images" / f"scene_{scene.selected_camera}_geometry_hdf5" required = ( detail / "metadata_scene.csv", detail / "metadata_cameras.csv", camera_detail / "camera_keyframe_frame_indices.hdf5", camera_detail / "camera_keyframe_look_at_positions.hdf5", camera_detail / "camera_keyframe_orientations.hdf5", camera_detail / "camera_keyframe_positions.hdf5", ) return ( all(path.is_file() and path.stat().st_size > 0 for path in required) and len(list(preview.glob("frame.*.color.jpg"))) == scene.frame_count and len(list(geometry.glob("frame.*.depth_meters.hdf5"))) == scene.frame_count and len(list(geometry.glob("frame.*.normal_world.hdf5"))) == scene.frame_count ) def _prepared_scene_complete(destination: Path, frame_count: int) -> bool: manifest_path = destination / "manifest.json" if not manifest_path.is_file(): return False try: return len(load_manifest(manifest_path).frames) == frame_count except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError): return False def _require_space(path: Path, reserve_bytes: int) -> None: free = shutil.disk_usage(path.resolve().anchor).free if free < reserve_bytes: raise OSError(f"free-space reserve violated: {free} < {reserve_bytes} bytes") def _append_jsonl(path: Path, record: dict[str, object]) -> None: with path.open("a", encoding="utf-8") as handle: handle.write(json.dumps(record, sort_keys=True) + "\n") if __name__ == "__main__": raise SystemExit(main())