scripts/download_pointworld.py

#!/usr/bin/env python3
"""CLI: selectively download PointWorld-DROID assets (manifest + cameras + one shard).

Never touches ``droid/depth_320x180/`` -- it is 1.23 TB and not
cherry-pickable; ``PointWorldStore`` refuses that prefix outright as a second
line of defense even if this script's flags were ever extended to reach it.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from fpgm.data.pointworld import HF_REPO_ID, PointWorldStore  # noqa: E402
from fpgm.utils.io import human_bytes  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger(__name__)

# Approximate *packed* sizes from the verified dataset listing; used only for
# the pre-flight confirmation gate below, never for validating a download.
_APPROX_BYTES = {
    "manifest": 3_900_000,
    "cameras": 19_500_000,
    "shard": 3_500_000_000,
}
_CONFIRM_THRESHOLD_BYTES = 1_000_000_000


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path, default=Path("data/pointworld"))
    parser.add_argument("--repo", default=HF_REPO_ID)
    parser.add_argument(
        "--shard",
        default="shard-000000",
        help="shard id to download and restore, or '' to skip shard download entirely",
    )
    parser.add_argument(
        "--skip-cameras", action="store_true", help="skip the *_cameras.json bundle"
    )
    parser.add_argument("--skip-manifest", action="store_true", help="skip _shards_manifest.json")
    parser.add_argument(
        "--yes", "-y", action="store_true", help="skip the >1 GB size confirmation gate"
    )
    parser.add_argument("--log-level", default="INFO")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    setup_logging(args.log_level)

    total_bytes = 0
    if not args.skip_manifest:
        total_bytes += _APPROX_BYTES["manifest"]
    if not args.skip_cameras:
        total_bytes += _APPROX_BYTES["cameras"]
    if args.shard:
        total_bytes += _APPROX_BYTES["shard"]

    print(f"About to download ~{human_bytes(total_bytes)} from {args.repo}", file=sys.stderr)
    if total_bytes > _CONFIRM_THRESHOLD_BYTES and not args.yes:
        print("This exceeds ~1 GB; re-run with --yes to proceed.", file=sys.stderr)
        return 1

    store = PointWorldStore(local_dir=args.data_dir, hf_repo=args.repo)

    if not args.skip_manifest:
        path = store.download_manifest()
        logger.info("manifest -> %s", path)
    if not args.skip_cameras:
        path = store.download_cameras()
        logger.info("cameras -> %s", path)
    if args.shard:
        path = store.download_shard(args.shard)
        episodes = store.episodes_in_shard(args.shard)
        logger.info("shard %s -> %s (%d episodes)", args.shard, path, len(episodes))

    return 0


if __name__ == "__main__":
    raise SystemExit(main())