Add missing batch GIF renderer

code/scripts/render_oven_gif_batch.py

@@ -0,0 +1,228 @@
+from pathlib import Path
+import argparse
+import json
+import os
+import shutil
+import subprocess
+import sys
+import time
+from typing import Dict, List, Optional, Sequence
+
+
+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",
+        "MALLOC_ARENA_MAX": "2",
+    }
+    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()
+
+from rr_label_study.oven_study import _episode_dirs
+
+
+def _select_episode_indices(
+    total_episodes: int,
+    episode_offset: int,
+    max_episodes: Optional[int],
+    episode_indices: Optional[Sequence[int]],
+) -> List[int]:
+    if episode_indices is not None:
+        selected: List[int] = []
+        seen = set()
+        for raw_index in episode_indices:
+            episode_index = int(raw_index)
+            if not (0 <= episode_index < total_episodes):
+                raise ValueError(
+                    f"episode index {episode_index} outside available range 0..{total_episodes - 1}"
+                )
+            if episode_index in seen:
+                continue
+            selected.append(episode_index)
+            seen.add(episode_index)
+        return selected
+
+    remaining = max(0, total_episodes - episode_offset)
+    if max_episodes is not None:
+        remaining = min(remaining, max_episodes)
+    if remaining <= 0:
+        return []
+    return list(range(episode_offset, episode_offset + remaining))
+
+
+def _has_gif_suite(episode_output_dir: Path, episode_name: str) -> bool:
+    visualizations_dir = episode_output_dir.joinpath("visualizations")
+    required = [
+        visualizations_dir.joinpath(f"{episode_name}_all_metrics.gif"),
+        visualizations_dir.joinpath(f"{episode_name}_visibility_focus.gif"),
+        visualizations_dir.joinpath(f"{episode_name}_path_quality_focus.gif"),
+        visualizations_dir.joinpath("README.md"),
+    ]
+    return all(path.exists() for path in required)
+
+
+def _write_json(path: Path, payload: Dict[str, object]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset-root",
+        default="/workspace/data/bimanual_take_tray_out_of_oven_train_128",
+    )
+    parser.add_argument("--result-dir", required=True)
+    parser.add_argument("--episode-offset", type=int, default=0)
+    parser.add_argument("--max-episodes", type=int, default=100)
+    parser.add_argument("--episode-indices")
+    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=1200)
+    args = parser.parse_args()
+
+    dataset_root = Path(args.dataset_root)
+    result_dir = Path(args.result_dir)
+    logs_dir = result_dir.joinpath("render_logs")
+    logs_dir.mkdir(parents=True, exist_ok=True)
+    progress_path = result_dir.joinpath("render_progress.json")
+
+    all_episode_dirs = _episode_dirs(dataset_root)
+    explicit_episode_indices = None
+    if args.episode_indices:
+        explicit_episode_indices = [
+            int(chunk.strip()) for chunk in args.episode_indices.split(",") if chunk.strip()
+        ]
+    selected_episode_indices = _select_episode_indices(
+        total_episodes=len(all_episode_dirs),
+        episode_offset=args.episode_offset,
+        max_episodes=args.max_episodes,
+        episode_indices=explicit_episode_indices,
+    )
+
+    completed: List[int] = []
+    for episode_index in selected_episode_indices:
+        episode_name = f"episode{episode_index}"
+        episode_output_dir = result_dir.joinpath(episode_name)
+        if not episode_output_dir.exists():
+            raise FileNotFoundError(f"missing episode output dir: {episode_output_dir}")
+        if _has_gif_suite(episode_output_dir, episode_name):
+            completed.append(episode_index)
+            _write_json(
+                progress_path,
+                {
+                    "current_episode": None,
+                    "completed_episode_indices": completed,
+                    "total_selected": len(selected_episode_indices),
+                    "updated_at_epoch": time.time(),
+                },
+            )
+            continue
+
+        temp_output_dir = episode_output_dir.joinpath("visualizations_tmp")
+        if temp_output_dir.exists():
+            shutil.rmtree(temp_output_dir)
+
+        log_path = logs_dir.joinpath(f"{episode_name}.log")
+        env = os.environ.copy()
+        env["OMP_NUM_THREADS"] = "1"
+        env["OPENBLAS_NUM_THREADS"] = "1"
+        env["MKL_NUM_THREADS"] = "1"
+        env["NUMEXPR_NUM_THREADS"] = "1"
+        env["VECLIB_MAXIMUM_THREADS"] = "1"
+        env["BLIS_NUM_THREADS"] = "1"
+        env["MALLOC_ARENA_MAX"] = "2"
+        env["PYTHONUNBUFFERED"] = "1"
+        _write_json(
+            progress_path,
+            {
+                "current_episode": episode_name,
+                "completed_episode_indices": completed,
+                "total_selected": len(selected_episode_indices),
+                "updated_at_epoch": time.time(),
+            },
+        )
+        with log_path.open("w", encoding="utf-8") as log_handle:
+            process = subprocess.Popen(
+                [
+                    sys.executable,
+                    str(PROJECT_ROOT.joinpath("scripts", "render_oven_metric_gifs.py")),
+                    "--episode-dir",
+                    str(all_episode_dirs[episode_index]),
+                    "--dense-csv",
+                    str(episode_output_dir.joinpath(f"{episode_name}.dense.csv")),
+                    "--templates-pkl",
+                    str(episode_output_dir.joinpath("templates.pkl")),
+                    "--output-dir",
+                    str(temp_output_dir),
+                    "--debug-jsonl",
+                    str(episode_output_dir.joinpath(f"{episode_name}.debug.jsonl")),
+                    "--checkpoint-stride",
+                    str(args.checkpoint_stride),
+                    "--num-workers",
+                    str(args.num_workers),
+                    "--base-display",
+                    str(args.base_display),
+                ],
+                stdout=log_handle,
+                stderr=subprocess.STDOUT,
+                cwd=str(PROJECT_ROOT),
+                env=env,
+            )
+            return_code = process.wait()
+        if return_code != 0:
+            raise RuntimeError(f"gif render failed for {episode_name}; see {log_path}")
+
+        final_visualizations_dir = episode_output_dir.joinpath("visualizations")
+        if final_visualizations_dir.exists():
+            shutil.rmtree(final_visualizations_dir)
+        temp_output_dir.rename(final_visualizations_dir)
+        completed.append(episode_index)
+        _write_json(
+            progress_path,
+            {
+                "current_episode": None,
+                "completed_episode_indices": completed,
+                "total_selected": len(selected_episode_indices),
+                "updated_at_epoch": time.time(),
+            },
+        )
+
+    _write_json(
+        progress_path,
+        {
+            "current_episode": None,
+            "completed_episode_indices": completed,
+            "total_selected": len(selected_episode_indices),
+            "finished_at_epoch": time.time(),
+        },
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())