scripts/data/render_trex_track_samples.py

#!/usr/bin/env python3
"""Render uniformly sampled T-Rex track-cache episodes without rerunning models."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

_SCRIPT_DIR = Path(__file__).resolve().parent
_REPO_ROOT = _SCRIPT_DIR.parents[1]
_SCRIPTS_DIR = _REPO_ROOT / "scripts"
if str(_SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(_SCRIPTS_DIR))

from extract_track import load_episode_videos  # noqa: E402
from trex_track.layout import VIEW_ORDER  # noqa: E402
from trex_track.trex_viz_tracks import render_three_view_combined_video  # noqa: E402


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dataset-root",
        type=Path,
        default=_REPO_ROOT / "data" / "trex_full_force",
    )
    parser.add_argument("--track-cache", type=Path, default=None)
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=_REPO_ROOT / "outputs" / "trex_track_samples",
    )
    parser.add_argument("--num-samples", type=int, default=3)
    parser.add_argument("--episodes", type=int, nargs="*", default=None)
    parser.add_argument("--fps", type=int, default=0)
    parser.add_argument("--trail", type=int, default=15)
    return parser.parse_args()


def _uniform_episode_indices(total_episodes: int, count: int) -> list[int]:
    if count < 1:
        raise ValueError("--num-samples must be positive")
    if count > total_episodes:
        raise ValueError("--num-samples cannot exceed total episodes")
    return [
        int(index)
        for index in np.rint(
            np.linspace(0, total_episodes - 1, count, dtype=np.float64)
        )
    ]


def _load_tracks(path: Path) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
    if not path.is_file():
        raise FileNotFoundError(path)
    with np.load(path, allow_pickle=False) as payload:
        tracks = {
            view: np.asarray(payload[f"tracks_{view}"], dtype=np.float32)
            for view in VIEW_ORDER
        }
        visibility = {
            view: np.asarray(payload[f"vis_{view}"], dtype=np.float32)
            for view in VIEW_ORDER
        }
    for view in VIEW_ORDER:
        if visibility[view].shape != tracks[view].shape[:2]:
            raise ValueError(
                f"{path}: {view} visibility {visibility[view].shape} "
                f"does not match tracks {tracks[view].shape}"
            )
    return tracks, visibility


def main() -> int:
    args = _parse_args()
    dataset_root = args.dataset_root.expanduser().resolve()
    track_cache = (
        args.track_cache.expanduser().resolve()
        if args.track_cache is not None
        else dataset_root / "tracks_trex_track_force_v2"
    )
    output_dir = args.output_dir.expanduser().resolve()
    info = json.loads((dataset_root / "meta" / "info.json").read_text())
    total_episodes = int(info["total_episodes"])
    video_shape = info["features"]["observation.images.head_left"]["shape"]
    out_hw = (int(video_shape[0]), int(video_shape[1]))
    fps = int(args.fps) if args.fps > 0 else int(info["fps"])

    episode_indices = (
        [int(index) for index in args.episodes]
        if args.episodes
        else _uniform_episode_indices(total_episodes, int(args.num_samples))
    )
    if any(index < 0 or index >= total_episodes for index in episode_indices):
        raise ValueError(f"episode indices must be in [0, {total_episodes})")

    output_dir.mkdir(parents=True, exist_ok=True)
    rendered: list[dict[str, object]] = []
    for episode_index in episode_indices:
        cache_path = track_cache / f"episode_{episode_index:06d}.npz"
        output_path = output_dir / f"episode_{episode_index:06d}_tracks.mp4"
        print(f"Rendering episode {episode_index}: {output_path}", flush=True)
        tracks, visibility = _load_tracks(cache_path)
        videos = load_episode_videos(
            dataset_root,
            episode_index,
            out_hw=out_hw,
        )
        frame_counts = {
            view: int(videos[view].shape[0])
            for view in VIEW_ORDER
        }
        for view in VIEW_ORDER:
            if int(tracks[view].shape[0]) != frame_counts[view]:
                raise ValueError(
                    f"episode {episode_index} {view}: "
                    f"{tracks[view].shape[0]} track frames != "
                    f"{frame_counts[view]} video frames"
                )
        render_three_view_combined_video(
            view_images=videos,
            view_tracks=tracks,
            view_vis=visibility,
            out_path=output_path,
            fps=fps,
            draw_trail=int(args.trail),
            dim_low_vis=True,
        )
        rendered.append(
            {
                "episode_index": episode_index,
                "frames": min(frame_counts.values()),
                "fps": fps,
                "track_cache": str(cache_path),
                "video": str(output_path),
            }
        )
        del tracks, visibility, videos

    summary_path = output_dir / "samples.json"
    summary_path.write_text(
        json.dumps(
            {
                "dataset_root": str(dataset_root),
                "sampling": "uniform endpoints and midpoint",
                "episodes": rendered,
            },
            indent=2,
        )
        + "\n"
    )
    print(f"Wrote {summary_path}")
    return 0


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