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())