scripts/download_droid_episodes.py

#!/usr/bin/env python3
"""CLI: download DROID raw metadata + camera mp4s for a set of episodes.

Episodes can be named explicitly (``--uuids``) or taken from the front of an
already-downloaded PointWorld shard (``--from-shard``, see
``download_pointworld.py``) so the raw video and the scene-flow annotations
line up for the same episodes.
"""

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.droid_raw import DroidRawClient  # noqa: E402
from fpgm.data.pointworld import PointWorldStore  # noqa: E402
from fpgm.types import DataError, EpisodeNotFoundError  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger(__name__)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--uuids", nargs="+", help="explicit episode uuids to fetch")
    source.add_argument(
        "--from-shard", help="take the first --n episodes of this downloaded PointWorld shard id"
    )
    parser.add_argument("--n", type=int, default=5, help="episode count when using --from-shard")
    parser.add_argument("--pointworld-dir", type=Path, default=Path("data/pointworld"))
    parser.add_argument(
        "--cameras-dir", type=Path, default=Path("data/pointworld/droid/cameras"),
        help="PointWorld cameras.json dir; supplies each episode's real bucket path "
             "(1.52%% of uuids do not derive correctly) and its success/failure outcome",
    )
    parser.add_argument("--data-dir", type=Path, default=Path("data/droid_raw"))
    parser.add_argument(
        "--roles", nargs="+", default=["ext1", "ext2"], help="camera roles to download mp4s for"
    )
    parser.add_argument("--no-trajectory", action="store_true", help="skip trajectory.h5")
    parser.add_argument("--log-level", default="INFO")
    return parser.parse_args()


def resolve_uuids(args: argparse.Namespace) -> list[str]:
    if args.uuids:
        return args.uuids
    store = PointWorldStore(local_dir=args.pointworld_dir)
    return store.episodes_in_shard(args.from_shard)[: args.n]


def main() -> int:
    args = parse_args()
    setup_logging(args.log_level)
    uuids = resolve_uuids(args)
    logger.info("downloading %d episode(s) -> %s", len(uuids), args.data_dir)

    client = DroidRawClient(cache_dir=args.data_dir, cameras_dir=args.cameras_dir)
    failures: list[str] = []
    for uuid in uuids:
        try:
            assets = client.download_episode(
                uuid, roles=tuple(args.roles), include_trajectory=not args.no_trajectory
            )
            logger.info(
                "%s: task=%r cameras=%s",
                uuid,
                assets.metadata.get("current_task"),
                sorted(assets.mp4_paths),
            )
        except (EpisodeNotFoundError, DataError) as exc:
            logger.error("%s: %s", uuid, exc)
            failures.append(uuid)

    if failures:
        logger.error("%d/%d episode(s) failed: %s", len(failures), len(uuids), failures)
        return 1
    return 0


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