scripts/export_vace_samples.py

#!/usr/bin/env python
"""S8 CLI: re-export the ``vace/`` training tier from an episode's ``master/`` artifacts.

A thin wrapper over :func:`fpgm.datagen.pipeline.run_s8_export` -- the same function
:class:`~fpgm.datagen.pipeline.EpisodePipeline`'s own S8 stage calls in-process. This
script exists only for the ``--caption`` override (a manual one-off render with a
hand-written caption), which :meth:`EpisodePipeline.run`'s fixed signature has no room
for; every other code path is identical to running ``run_datagen.py --stages s8``.

Requires S2 (robot buffers + chosen extrinsics), S3 (``background_depth.h5``/
``plate.png``) and S6 (``poses.npz`` + ``master/meshes/*``) to already exist for this
uuid/camera.

Usage:
    PYTHONPATH=src python scripts/export_vace_samples.py \\
        --uuid AUTOLab+0d4edc83+2023-10-21-19h-07m-04s --camera 22008760
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

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

from fpgm.config_datagen import DatagenProfile  # noqa: E402
from fpgm.datagen.episode_spec import EpisodeMetadata, EpisodeSpecResolver  # noqa: E402
from fpgm.datagen.pipeline import run_s8_export  # noqa: E402
from fpgm.types import DataError  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger("export_vace_samples")

_DEFAULT_CONFIG = REPO_ROOT / "configs" / "datagen_droid.yaml"


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--uuid", required=True)
    p.add_argument(
        "--camera", required=True, help="camera serial (e.g. 22008760) or role (ext1/ext2)"
    )
    p.add_argument("--config", type=Path, default=_DEFAULT_CONFIG, help="DatagenProfile YAML")
    p.add_argument(
        "--out-root", type=Path, default=None,
        help="override the profile's paths.output_root",
    )
    p.add_argument("--caption", default=None, help="override the DROID current_task caption")
    p.add_argument("--force", action="store_true", help="ignore the on-disk window cache")
    p.add_argument("--log-level", default="INFO")
    return p.parse_args()


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

    profile = DatagenProfile.from_yaml(args.config)
    if args.out_root is not None:
        profile.paths.output_root = args.out_root.resolve()

    metadata = EpisodeMetadata.from_flows_h5(profile.paths.flows_h5(args.uuid))
    spec = EpisodeSpecResolver(profile).resolve(metadata, camera_role=args.camera)

    master_dir = profile.paths.master_dir(spec.uuid, spec.camera_serial)
    if not master_dir.exists():
        raise DataError(
            f"no master/ directory at {master_dir} -- run run_datagen.py --stages s6 first"
        )

    payload = run_s8_export(
        profile, spec, master_dir, force=args.force, caption_override=args.caption
    )

    caption_kind = "template-composed" if payload["caption_is_template_composed"] else "override"
    print(f"\ncaption ({caption_kind}): {payload['caption']!r}")
    print(f"exported {payload['n_windows']} window(s):")
    for w in payload["windows"]:
        res = f"{w['resolution'][0]}x{w['resolution'][1]}"
        print(
            f"  [{w['video_frame_start']:4d}, {w['video_frame_end']:4d})  {res}  "
            f"overlaps_pose_gap={w['overlaps_pose_gap']}  -> {Path(w['target_video']).parent}"
        )
    return 0


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