code/scripts/render_oven_metric_gifs.py

from pathlib import Path
import argparse
import os
import signal
import subprocess
import sys
import time
from typing import Dict, List, Optional, Tuple

from PIL import Image
import numpy as np
import pandas as pd


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def _configure_thread_env() -> None:
    defaults = {
        "OMP_NUM_THREADS": "1",
        "OPENBLAS_NUM_THREADS": "1",
        "MKL_NUM_THREADS": "1",
        "NUMEXPR_NUM_THREADS": "1",
        "VECLIB_MAXIMUM_THREADS": "1",
        "BLIS_NUM_THREADS": "1",
    }
    for key, value in defaults.items():
        os.environ.setdefault(key, value)


def _configure_coppeliasim_env() -> None:
    coppeliasim_root = os.environ.setdefault("COPPELIASIM_ROOT", "/workspace/coppelia_sim")
    ld_library_path_parts = [
        part for part in os.environ.get("LD_LIBRARY_PATH", "").split(":") if part
    ]
    if coppeliasim_root not in ld_library_path_parts:
        ld_library_path_parts.insert(0, coppeliasim_root)
    os.environ["LD_LIBRARY_PATH"] = ":".join(ld_library_path_parts)


_configure_thread_env()
_configure_coppeliasim_env()


def _launch_xvfb(display_num: int, log_path: Path) -> subprocess.Popen:
    log_handle = log_path.open("w", encoding="utf-8")
    return subprocess.Popen(
        [
            "Xvfb",
            f":{display_num}",
            "-screen",
            "0",
            "1280x1024x24",
            "+extension",
            "GLX",
            "+render",
            "-noreset",
        ],
        stdout=log_handle,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )


def _stop_process(process: Optional[subprocess.Popen]) -> None:
    if process is None or process.poll() is not None:
        return
    try:
        os.killpg(process.pid, signal.SIGTERM)
    except ProcessLookupError:
        return
    try:
        process.wait(timeout=10)
    except subprocess.TimeoutExpired:
        try:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass


def _spawn_worker(
    display_num: int,
    episode_dir: Path,
    templates_pkl: Path,
    dense_csv: Path,
    debug_jsonl: Optional[Path],
    frame_index: int,
    checkpoint_stride: int,
    visibility_out: Optional[Path],
    path_out: Optional[Path],
    all_out: Optional[Path],
) -> subprocess.Popen:
    runtime_dir = Path(f"/tmp/rr_metric_gifs_display_{display_num}")
    runtime_dir.mkdir(parents=True, exist_ok=True)
    env = os.environ.copy()
    env["DISPLAY"] = f":{display_num}"
    env["COPPELIASIM_ROOT"] = "/workspace/coppelia_sim"
    env["LD_LIBRARY_PATH"] = f"/workspace/coppelia_sim:{env.get('LD_LIBRARY_PATH', '')}"
    env["QT_QPA_PLATFORM_PLUGIN_PATH"] = "/workspace/coppelia_sim"
    env["XDG_RUNTIME_DIR"] = str(runtime_dir)
    command = [
        sys.executable,
        str(PROJECT_ROOT.joinpath("scripts", "render_oven_metric_frame.py")),
        "--episode-dir",
        str(episode_dir),
        "--templates-pkl",
        str(templates_pkl),
        "--dense-csv",
        str(dense_csv),
        "--frame-index",
        str(frame_index),
        "--checkpoint-stride",
        str(checkpoint_stride),
    ]
    if debug_jsonl is not None:
        command.extend(["--debug-jsonl", str(debug_jsonl)])
    if visibility_out is not None:
        command.extend(["--visibility-out", str(visibility_out)])
    if path_out is not None:
        command.extend(["--path-out", str(path_out)])
    if all_out is not None:
        command.extend(["--all-out", str(all_out)])
    return subprocess.Popen(
        command,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        cwd=str(PROJECT_ROOT),
        env=env,
        start_new_session=True,
    )


def _assemble_gif(frame_paths: List[Path], output_path: Path, duration_ms: int) -> None:
    images = [Image.open(path).convert("P", palette=Image.Palette.ADAPTIVE) for path in frame_paths]
    output_path.parent.mkdir(parents=True, exist_ok=True)
    images[0].save(
        output_path,
        save_all=True,
        append_images=images[1:],
        duration=duration_ms,
        loop=0,
        disposal=2,
    )


def _range_inclusive(start: int, end: int, step: int) -> List[int]:
    if end < start:
        return []
    return list(range(start, end + 1, step))


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--episode-dir",
        default="/workspace/data/bimanual_take_tray_out_of_oven_train_128/all_variations/episodes/episode0",
    )
    parser.add_argument(
        "--dense-csv",
        default="/workspace/reveal_retrieve_label_study/results/oven_episode0_repaired_v1/episode0.dense.csv",
    )
    parser.add_argument(
        "--templates-pkl",
        default="/workspace/reveal_retrieve_label_study/results/oven_episode0_repaired_v1/templates.pkl",
    )
    parser.add_argument(
        "--output-dir",
        default="/workspace/reveal_retrieve_label_study/results/oven_episode0_repaired_v1/visualizations",
    )
    parser.add_argument("--debug-jsonl")
    parser.add_argument("--checkpoint-stride", type=int, default=16)
    parser.add_argument("--num-workers", type=int, default=6)
    parser.add_argument("--base-display", type=int, default=190)
    parser.add_argument("--all-metrics-only", action="store_true")
    args = parser.parse_args()

    episode_dir = Path(args.episode_dir)
    dense_csv = Path(args.dense_csv)
    templates_pkl = Path(args.templates_pkl)
    output_dir = Path(args.output_dir)
    debug_jsonl = Path(args.debug_jsonl) if args.debug_jsonl else dense_csv.with_name(
        dense_csv.name.replace(".dense.csv", ".debug.jsonl")
    )
    if not debug_jsonl.exists():
        debug_jsonl = None
    frame_df = pd.read_csv(dense_csv)
    episode_name = episode_dir.name
    frame_indices = (
        frame_df["frame_index"].to_numpy(dtype=int)
        if "frame_index" in frame_df
        else np.arange(len(frame_df), dtype=int)
    )
    max_frame = int(frame_indices.max()) if len(frame_indices) else 0
    phase_candidates = frame_df.loc[frame_df["phase_switch"].to_numpy(dtype=float) >= 0.5]
    if len(phase_candidates):
        phase_cross = int(phase_candidates.iloc[0]["frame_index"])
    else:
        ppre_candidates = frame_df.loc[frame_df["p_pre"].to_numpy(dtype=float) >= 0.45]
        phase_cross = int(ppre_candidates.iloc[0]["frame_index"]) if len(ppre_candidates) else max_frame // 2
    pext_candidates = frame_df.loc[frame_df["p_ext"].to_numpy(dtype=float) >= 0.45]
    pext_cross = int(pext_candidates.iloc[0]["frame_index"]) if len(pext_candidates) else phase_cross
    ready_candidates = frame_df.loc[frame_df["y_ready"].to_numpy(dtype=float) >= 0.5]
    ready_cross = int(ready_candidates.iloc[0]["frame_index"]) if len(ready_candidates) else pext_cross
    retrieve_candidates = frame_df.loc[frame_df["y_retrieve"].to_numpy(dtype=float) >= 0.5]
    retrieve_cross = int(retrieve_candidates.iloc[0]["frame_index"]) if len(retrieve_candidates) else max(phase_cross, pext_cross)

    frames_dir = output_dir.joinpath("frames")
    visibility_dir = frames_dir.joinpath("visibility_focus")
    path_dir = frames_dir.joinpath("path_quality_focus")
    all_dir = frames_dir.joinpath("all_metrics")
    if args.all_metrics_only:
        all_dir.mkdir(parents=True, exist_ok=True)
    else:
        for directory in [visibility_dir, path_dir, all_dir]:
            directory.mkdir(parents=True, exist_ok=True)

    if args.all_metrics_only:
        visibility_frames = []
        path_frames = []
    else:
        visibility_frames = _range_inclusive(
            max(0, phase_cross - 12), min(max_frame, phase_cross + 28), 1
        )
        path_start = max(0, min(phase_cross, pext_cross) - 16)
        path_end = min(max_frame, max(pext_cross, ready_cross, retrieve_cross) + 24)
        path_frames = sorted(
            set(
                _range_inclusive(
                    path_start, min(path_end, max(path_start, pext_cross - 18)), 2
                )
                + _range_inclusive(max(path_start, pext_cross - 18), path_end, 1)
            )
        )
    all_frames = _range_inclusive(0, max_frame, 2)
    unique_frames = sorted(set(visibility_frames) | set(path_frames) | set(all_frames))

    pending = unique_frames[:]
    active: Dict[int, Tuple[int, subprocess.Popen]] = {}
    displays = [args.base_display + i for i in range(args.num_workers)]
    xvfb_procs: List[subprocess.Popen] = []
    try:
        for display_num in displays:
            xvfb_procs.append(_launch_xvfb(display_num, output_dir.joinpath(f"xvfb_{display_num}.log")))
        time.sleep(1.0)

        while pending or active:
            free_displays = [display for display in displays if display not in active]
            while pending and free_displays:
                display_num = free_displays.pop(0)
                frame_index = pending.pop(0)
                process = _spawn_worker(
                    display_num=display_num,
                    episode_dir=episode_dir,
                    templates_pkl=templates_pkl,
                    dense_csv=dense_csv,
                    debug_jsonl=debug_jsonl,
                    frame_index=frame_index,
                    checkpoint_stride=args.checkpoint_stride,
                    visibility_out=visibility_dir.joinpath(f"frame_{frame_index:04d}.png")
                    if frame_index in visibility_frames
                    else None,
                    path_out=path_dir.joinpath(f"frame_{frame_index:04d}.png")
                    if frame_index in path_frames
                    else None,
                    all_out=all_dir.joinpath(f"frame_{frame_index:04d}.png")
                    if frame_index in all_frames
                    else None,
                )
                active[display_num] = (frame_index, process)
            time.sleep(0.5)
            finished = []
            for display_num, (frame_index, process) in active.items():
                return_code = process.poll()
                if return_code is None:
                    continue
                if return_code != 0:
                    raise RuntimeError(f"render worker failed for frame {frame_index} on display :{display_num}")
                finished.append(display_num)
            for display_num in finished:
                active.pop(display_num)
    finally:
        for _, process in list(active.values()):
            _stop_process(process)
        for xvfb in xvfb_procs:
            _stop_process(xvfb)

    visibility_pngs = [visibility_dir.joinpath(f"frame_{frame:04d}.png") for frame in visibility_frames]
    path_pngs = [path_dir.joinpath(f"frame_{frame:04d}.png") for frame in path_frames]
    all_pngs = [all_dir.joinpath(f"frame_{frame:04d}.png") for frame in all_frames]
    for path in all_pngs + visibility_pngs + path_pngs:
        if not path.exists():
            raise RuntimeError(f"missing rendered frame: {path}")

    if not args.all_metrics_only:
        _assemble_gif(
            visibility_pngs,
            output_dir.joinpath(f"{episode_name}_visibility_focus.gif"),
            duration_ms=120,
        )
        _assemble_gif(
            path_pngs,
            output_dir.joinpath(f"{episode_name}_path_quality_focus.gif"),
            duration_ms=160,
        )
    _assemble_gif(
        all_pngs,
        output_dir.joinpath(f"{episode_name}_all_metrics.gif"),
        duration_ms=100,
    )

    readme_lines = [
        "# Visualizations",
        "",
        f"- `{episode_name}_all_metrics.gif`: full episode overlay GIF over dense frames 0-{max_frame}, sampled every 2 frames.",
        "- `frames/all_metrics/`: per-frame PNGs with the episode-wide metric bars and phase banner.",
    ]
    if not args.all_metrics_only:
        readme_lines.extend(
            [
                f"- `{episode_name}_visibility_focus.gif`: three-view visibility montage over dense frames {visibility_frames[0]}-{visibility_frames[-1]}.",
                f"- `{episode_name}_path_quality_focus.gif`: debug-aware p_ext planner montage over dense frames {path_frames[0]}-{path_frames[-1]}.",
                "- `frames/visibility_focus/`: per-frame PNGs with scene overlays, x-ray projections, and tray-mask views.",
                "- `frames/path_quality_focus/`: per-frame PNGs with demo wrist trails, milestone pose overlays, and p_ext planner-search tables from the debug JSONL sidecar.",
            ]
        )
    readme_lines.extend(
        [
            "",
            "Legend highlights:",
            "- All metrics: red banner = reveal phase, green banner = retrieve phase.",
        ]
    )
    if not args.all_metrics_only:
        readme_lines.extend(
            [
                "- Visibility: blue = sampled tray surface, magenta = sampled grasp region, green = depth-consistent visible grasp samples.",
                "- Path quality: cyan/purple = recent left/right demo wrist trails, yellow = pregrasp, orange = grasp, green = retreat / extraction path.",
            ]
        )
    readme = output_dir.joinpath("README.md")
    readme.write_text("\n".join(readme_lines), encoding="utf-8")
    print(output_dir)
    return 0


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