scripts/run_pipeline.py

#!/usr/bin/env python
"""Run the velocity-extraction pipeline over a set of DROID episodes.

    # IMPORTANT: the GPU driver shim must be active first (see README)
    source scripts/nvidia_lib_shim.sh
    conda activate fpgm
    python scripts/run_pipeline.py --config configs/droid_velocity.yaml

Results are written under ``outputs/<episode-uuid>/`` and a JSON summary is printed
so a batch run is inspectable without opening the videos.
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import asdict, is_dataclass
from pathlib import Path

import numpy as np

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

from fpgm.config import PipelineConfig  # noqa: E402
from fpgm.pipeline.velocity_pipeline import EpisodeResult, VelocityPipeline  # noqa: E402
from fpgm.utils.gpu import ensure_cuda, log_gpu_memory  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger("run_pipeline")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument(
        "--config",
        type=Path,
        default=REPO_ROOT / "configs" / "droid_velocity.yaml",
        help="pipeline config YAML",
    )
    p.add_argument(
        "--episodes",
        nargs="*",
        default=None,
        help="episode uuids to process (overrides the config list)",
    )
    p.add_argument(
        "--camera-role",
        default=None,
        choices=["ext1", "ext2"],
        help="which exterior camera to use (wrist has no scene-flow annotations)",
    )
    p.add_argument("--device", default=None, help="cuda | cuda:N | cpu")
    p.add_argument(
        "--summary-json",
        type=Path,
        default=None,
        help="also write the run summary to this path",
    )
    p.add_argument(
        "--skip-gpu-check",
        action="store_true",
        help="do not fail fast when CUDA is unavailable (CPU is impractically slow)",
    )
    return p.parse_args()


def summarise(results: list[EpisodeResult]) -> dict:
    """Build a JSON-serialisable digest of a run."""
    episodes = []
    for ep in results:
        clips = []
        for c in ep.clips:
            entry: dict = {
                "clip": c.clip_key,
                "camera": c.camera_serial,
                "prompt": c.prompt,
                "obj_id": c.obj_id,
                "convention": c.convention.value if c.convention else None,
                "ok": c.ok,
                "error": c.error,
                "depth_coverage": round(c.depth_coverage, 4),
                "depth_confidence": round(c.depth_confidence, 4),
                "artifacts": {k: str(v) for k, v in c.artifacts.items()},
            }
            if c.velocity is not None:
                speed = c.velocity.object_speed
                finite = np.isfinite(speed)
                entry["velocity"] = {
                    "frame": c.velocity.frame,
                    "frames_resolved": int(finite.sum()),
                    "frames_total": int(finite.size),
                    "peak_speed_mps": float(np.nanmax(speed)) if finite.any() else None,
                    "mean_speed_mps": float(np.nanmean(speed)) if finite.any() else None,
                }
            if c.validation is not None and is_dataclass(c.validation):
                entry["validation"] = _jsonable(asdict(c.validation))
            clips.append(entry)
        episodes.append(
            {
                "uuid": ep.episode_uuid,
                "task": ep.task_instruction,
                "ok": ep.ok,
                "error": ep.error,
                "clips": clips,
            }
        )

    n_ok = sum(1 for e in results if e.ok)
    return {
        "episodes_total": len(results),
        "episodes_ok": n_ok,
        "episodes": episodes,
    }


def _jsonable(obj):
    """Make numpy types and arrays JSON-safe, summarising large arrays."""
    if isinstance(obj, dict):
        return {k: _jsonable(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_jsonable(v) for v in obj]
    if isinstance(obj, np.ndarray):
        if obj.size > 32:
            finite = obj[np.isfinite(obj)] if obj.dtype.kind == "f" else obj
            return {
                "shape": list(obj.shape),
                "mean": float(finite.mean()) if finite.size else None,
                "median": float(np.median(finite)) if finite.size else None,
            }
        return _jsonable(obj.tolist())
    if isinstance(obj, (np.floating, np.integer)):
        return obj.item()
    if isinstance(obj, Path):
        return str(obj)
    return obj


def main() -> int:
    args = parse_args()
    cfg = PipelineConfig.from_yaml(args.config)

    if args.camera_role:
        cfg.dataset.camera_role = args.camera_role
    if args.device:
        cfg.runtime.device = args.device

    setup_logging(cfg.runtime.log_level)
    logger.info("config: %s", args.config)

    if not args.skip_gpu_check and cfg.runtime.device.startswith("cuda"):
        # Fail fast with the shim remedy rather than dying deep inside a model build.
        ensure_cuda()
        log_gpu_memory()

    with VelocityPipeline(cfg) as pipeline:
        results = pipeline.run(args.episodes)

    print("\n" + "=" * 78)
    for ep in results:
        for clip in ep.clips:
            print(clip.summary())
    print("=" * 78)

    summary = summarise(results)
    out_path = args.summary_json or (cfg.paths.outputs_dir / "run_summary.json")
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(summary, indent=2))
    logger.info("summary written to %s", out_path)

    ok = summary["episodes_ok"]
    print(f"\n{ok}/{summary['episodes_total']} episodes produced velocities")
    return 0 if ok else 1


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