""" Side-by-side viewer for forward-dynamics rollouts vs. ground truth. Point it at any HF dataset repo laid out as: windows_manifest.json # optional: list of window dicts metrics_summary.json # optional: per-window MSE / PSNR windows//gt.mp4 # ground-truth window windows//generated.mp4 # model rollout for the same slice Ground truth and generated are composed into a single video so the two panels stay frame-locked; two independent players drift apart and cannot be scrubbed together. Videos may stack N cameras vertically; N is inferred from the frame aspect ratio and the camera selector adapts. The app starts empty and assumes no dataset: type a repo id and press Load. Nothing is downloaded at build time and no repo is ever pulled whole. Each browser session gets its own scratch directory, keeps at most MAX_WINDOWS_ON_DISK windows, and is deleted when the session ends. Private repos need an HF_TOKEN secret with read access. """ from __future__ import annotations import json import os import re import shutil import tempfile import threading import time import uuid from collections import OrderedDict from pathlib import Path import cv2 import gradio as gr import imageio_ffmpeg import numpy as np from huggingface_hub import HfApi, RepoFolder, hf_hub_download HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None BASE_DIR = Path(tempfile.gettempdir()) / "wmviz" BASE_DIR.mkdir(parents=True, exist_ok=True) MAX_WINDOWS_ON_DISK = int(os.environ.get("MAX_WINDOWS_ON_DISK", 12)) MAX_PAIRS_IN_RAM = 2 MAX_COMPOSED_IN_RAM = 1 MAX_DISCOVER = int(os.environ.get("MAX_DISCOVER", 5000)) SESSION_TTL = int(os.environ.get("SESSION_TTL", 3600)) THREE_CAM_NAMES = ["Top", "Left wrist", "Right wrist"] ALL_CAMS = -1 NO_CRIT = -1 FONT = cv2.FONT_HERSHEY_SIMPLEX HDR, FTR, SEP = 26, 26, 4 SCALE = 2 # source frames are small; upscale so overlays stay legible GT_COLOR = (130, 225, 140) GEN_COLOR = (120, 180, 255) DIFF_COLOR = (215, 215, 215) CRIT_COLOR = (255, 95, 95) SORTS = [ "Hardest first (critical PSNR up)", "Easiest first (critical PSNR down)", "Task A-Z", "Manifest order", ] # --------------------------------------------------------------------------- session store def parse_repo(text: str) -> tuple[str, str | None]: """Split "owner/name@revision" into (repo_id, revision). Tolerates a pasted URL.""" text = (text or "").strip().rstrip("/") m = re.match(r"^https?://huggingface\.co/(?:datasets/)?(.+?)(?:/tree/([^/]+))?$", text) if m: return m.group(1), m.group(2) if "@" in text: repo, _, rev = text.partition("@") return repo.strip(), rev.strip() or None return text, None class Session: """One browser session's scratch space: bounded on disk and in RAM.""" def __init__(self, sid: str): self.sid = sid self.root = BASE_DIR / sid self.data = self.root / "data" self.render = self.root / "render" for d in (self.data, self.render): d.mkdir(parents=True, exist_ok=True) self.repo = "" self.revision: str | None = None self.rows: list[dict] = [] self._on_disk: OrderedDict[str, bool] = OrderedDict() self._pairs: OrderedDict[str, tuple[np.ndarray, np.ndarray]] = OrderedDict() self._composed: OrderedDict[tuple, np.ndarray] = OrderedDict() self.touched = time.time() def fetch(self, rel: str, optional: bool = False) -> str | None: """Download one repo file into this session's directory.""" self.touched = time.time() try: return hf_hub_download( repo_id=self.repo, filename=rel, repo_type="dataset", revision=self.revision, token=HF_TOKEN, local_dir=str(self.data), ) except Exception: if optional: return None raise def window_file(self, episode_id: str, name: str, optional: bool = False) -> str | None: path = self.fetch(f"windows/{episode_id}/{name}", optional=optional) if path: self._keep(episode_id) return path def _keep(self, episode_id: str) -> None: """Register a window as resident and evict the least recently used ones.""" self._on_disk[episode_id] = True self._on_disk.move_to_end(episode_id) while len(self._on_disk) > MAX_WINDOWS_ON_DISK: old, _ = self._on_disk.popitem(last=False) shutil.rmtree(self.data / "windows" / old, ignore_errors=True) self._pairs.pop(old, None) for key in [k for k in self._composed if k[0] == old]: self._composed.pop(key, None) for f in self.render.glob(f"{old}_*"): f.unlink(missing_ok=True) def disk_windows(self) -> int: return len(self._on_disk) def pair(self, episode_id: str) -> tuple[np.ndarray, np.ndarray]: hit = self._pairs.get(episode_id) if hit is not None: self._pairs.move_to_end(episode_id) return hit gt = read_video(self.window_file(episode_id, "gt.mp4")) gen = read_video(self.window_file(episode_id, "generated.mp4")) t = min(len(gt), len(gen)) value = (gt[:t], gen[:t]) self._pairs[episode_id] = value while len(self._pairs) > MAX_PAIRS_IN_RAM: self._pairs.popitem(last=False) return value def composed(self, episode_id: str, view: int, show_diff: bool, crit: int) -> np.ndarray: key = (episode_id, view, show_diff, crit) hit = self._composed.get(key) if hit is not None: self._composed.move_to_end(key) return hit value = compose(self.pair(episode_id), view, show_diff, crit) self._composed[key] = value while len(self._composed) > MAX_COMPOSED_IN_RAM: self._composed.popitem(last=False) return value def set_repo(self, repo: str, revision: str | None) -> None: if (repo, revision) != (self.repo, self.revision): self.reset() self.repo, self.revision = repo, revision def reset(self) -> None: self._on_disk.clear() self._pairs.clear() self._composed.clear() self.rows = [] for d in (self.data, self.render): shutil.rmtree(d, ignore_errors=True) d.mkdir(parents=True, exist_ok=True) def close(self) -> None: self._pairs.clear() self._composed.clear() shutil.rmtree(self.root, ignore_errors=True) SESSIONS: dict[str, Session] = {} _LOCK = threading.Lock() def sweep() -> None: """Drop sessions whose browser tab went away without a clean callback.""" now = time.time() with _LOCK: stale = [s for s in SESSIONS.values() if now - s.touched > SESSION_TTL] for s in stale: SESSIONS.pop(s.sid, None) s.close() live = set(SESSIONS) for d in BASE_DIR.iterdir(): # dirs left behind by a previous process try: if d.is_dir() and d.name not in live and now - d.stat().st_mtime > SESSION_TTL: shutil.rmtree(d, ignore_errors=True) except OSError: pass def get_session(sid: str | None) -> Session: with _LOCK: s = SESSIONS.get(sid or "") if s is None: sid = sid or uuid.uuid4().hex s = SESSIONS[sid] = Session(sid) s.touched = time.time() return s def release(sid: str | None) -> None: """gr.State delete_callback - fires when the session ends or its TTL expires.""" with _LOCK: s = SESSIONS.pop(sid or "", None) if s is not None: s.close() # --------------------------------------------------------------------------- index def _row(order: int, episode_id: str, w: dict, m: dict) -> dict: return { "order": order, "episode_id": episode_id, "task": w.get("task") or "?", "edge": w.get("edge_type") or "?", "arm": w.get("arm") or "?", "critical_frame": w.get("critical_frame"), "start_frame": w.get("start_frame"), "full_psnr": m.get("full_psnr"), "crit_psnr": m.get("crit_psnr"), "full_mse": m.get("full_mse"), "crit_mse": m.get("crit_mse"), } def discover_episodes(session: Session) -> tuple[list[str], bool]: """List windows/ one level deep, capped. Returns (ids, was_truncated).""" tree = HfApi().list_repo_tree( session.repo, path_in_repo="windows", recursive=False, revision=session.revision, repo_type="dataset", token=HF_TOKEN, ) eps, truncated = [], False for entry in tree: if not isinstance(entry, RepoFolder): continue eps.append(entry.path.split("/")[-1]) if len(eps) >= MAX_DISCOVER: truncated = True break return sorted(eps), truncated def build_index(session: Session) -> str: """Populate session.rows from the manifest, or by listing the repo. Returns a note.""" metrics = [] path = session.fetch("metrics_summary.json", optional=True) if path: try: metrics = json.load(open(path)) except Exception: metrics = [] by_id = {m.get("id"): m for m in metrics if isinstance(m, dict)} manifest = None path = session.fetch("windows_manifest.json", optional=True) if path: try: loaded = json.load(open(path)) if isinstance(loaded, list) and loaded: manifest = loaded except Exception: manifest = None if manifest is not None: session.rows = [ _row(i, w.get("episode_id", f"window_{i}"), w, by_id.get(w.get("episode_id"), {})) for i, w in enumerate(manifest) ] note = "`windows_manifest.json`" else: eps, truncated = discover_episodes(session) if not eps: raise FileNotFoundError( "no windows// directories found - is this the right repo?" ) session.rows = [_row(i, ep, {}, by_id.get(ep, {})) for i, ep in enumerate(eps)] note = "directory listing (no `windows_manifest.json`)" if truncated: note += f", capped at {MAX_DISCOVER}" if metrics: note += " + `metrics_summary.json`" return note def short_id(episode_id: str) -> str: return re.sub(r"^(episode|window|ep)[-_]", "", episode_id)[:8] def label_for(row: dict) -> str: psnr = row["crit_psnr"] head = f"{psnr:5.1f} dB" if isinstance(psnr, (int, float)) else " -- dB" tags = " ".join(t for t in (row["edge"], row["arm"]) if t and t != "?") parts = [head] + ([tags] if tags else []) + [row["task"], short_id(row["episode_id"])] return " | ".join(parts) def filter_rows(session: Session, query: str, sort: str) -> list[dict]: rows = session.rows q = (query or "").strip().lower() if q: rows = [r for r in rows if q in r["task"].lower() or q in r["episode_id"].lower()] def psnr(r, missing): v = r["crit_psnr"] return v if isinstance(v, (int, float)) else missing if sort == SORTS[0]: return sorted(rows, key=lambda r: psnr(r, float("inf"))) if sort == SORTS[1]: return sorted(rows, key=lambda r: -psnr(r, float("-inf"))) if sort == SORTS[2]: return sorted(rows, key=lambda r: (r["task"], r["episode_id"])) return sorted(rows, key=lambda r: r["order"]) def row_by_id(session: Session, episode_id: str) -> dict: for r in session.rows: if r["episode_id"] == episode_id: return r return _row(0, episode_id, {}, {}) # --------------------------------------------------------------------------- video def read_video(path: str) -> np.ndarray: """Decode an mp4 into a uint8 [T, H, W, 3] RGB array.""" cap = cv2.VideoCapture(path) frames = [] while True: ok, frame = cap.read() if not ok: break frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) cap.release() if not frames: raise RuntimeError(f"no frames decoded from {path}") return np.stack(frames) def n_cameras(frames: np.ndarray) -> int: """How many roughly-square cameras are stacked vertically in the frame.""" h, w = frames.shape[1], frames.shape[2] return max(1, min(6, int(round(h / w)))) def camera_names(n: int) -> list[str]: return THREE_CAM_NAMES if n == 3 else [f"Camera {i + 1}" for i in range(n)] def camera_choices(n: int) -> list[tuple[str, int]]: choices = [(name, i) for i, name in enumerate(camera_names(n))] return ([("All cameras", ALL_CAMS)] + choices) if n > 1 else choices def cam_slices(frames: np.ndarray, view: int) -> list[np.ndarray]: n = n_cameras(frames) h = frames.shape[1] // n rows = range(n) if view == ALL_CAMS or not 0 <= view < n else [view] return [frames[:, i * h : (i + 1) * h] for i in rows] def letterbox_rows(*bands: np.ndarray, thresh: int = 16) -> tuple[int, int]: """Rows to keep after dropping the black bars shared by every given band.""" h = bands[0].shape[1] lit = np.zeros(h, bool) for b in bands: lit |= b.reshape(b.shape[0], h, -1).max(axis=(0, 2)) > thresh idx = np.flatnonzero(lit) if idx.size < h * 0.1: # nothing meaningful to trim return 0, h return int(idx[0]), int(idx[-1]) + 1 def crop_view(gt: np.ndarray, gen: np.ndarray, view: int) -> tuple[np.ndarray, np.ndarray]: """Select cameras and trim the letterbox, identically for both videos.""" gt_out, gen_out = [], [] for g, p in zip(cam_slices(gt, view), cam_slices(gen, view)): lo, hi = letterbox_rows(g, p) gt_out.append(g[:, lo:hi]) gen_out.append(p[:, lo:hi]) return np.concatenate(gt_out, axis=1), np.concatenate(gen_out, axis=1) def upscale(frames: np.ndarray, scale: int) -> np.ndarray: if scale == 1: return frames return np.stack( [cv2.resize(f, None, fx=scale, fy=scale, interpolation=cv2.INTER_NEAREST) for f in frames] ) def panel(frames: np.ndarray, title: str, color: tuple[int, int, int]) -> np.ndarray: t, h, w, _ = frames.shape out = np.zeros((t, HDR + h, w, 3), np.uint8) out[:, HDR:] = frames bar = np.full((HDR, w, 3), 22, np.uint8) cv2.putText(bar, title, (8, HDR - 8), FONT, 0.52, color, 1, cv2.LINE_AA) out[:, :HDR] = bar return out def diff_frames(gt: np.ndarray, gen: np.ndarray, gain: float = 3.0) -> np.ndarray: d = np.abs(gt.astype(np.int16) - gen.astype(np.int16)).mean(axis=-1) d = np.clip(d * gain, 0, 255).astype(np.uint8) return np.stack([cv2.applyColorMap(f, cv2.COLORMAP_INFERNO)[:, :, ::-1] for f in d]) def compose(pair: tuple[np.ndarray, np.ndarray], view: int, show_diff: bool, crit_idx: int) -> np.ndarray: """Build the labelled side-by-side strip: [T, H, W, 3] RGB.""" gt_c, gen_c = crop_view(pair[0], pair[1], view) gt_v, gen_v = upscale(gt_c, SCALE), upscale(gen_c, SCALE) panels = [panel(gt_v, "GROUND TRUTH", GT_COLOR), panel(gen_v, "GENERATED", GEN_COLOR)] if show_diff: panels.append(panel(diff_frames(gt_v, gen_v), "|DIFF| x3", DIFF_COLOR)) t, ph, _, _ = panels[0].shape gap = np.zeros((t, ph, SEP, 3), np.uint8) stitched = [panels[0]] for p in panels[1:]: stitched += [gap, p] comp = np.concatenate(stitched, axis=2) t, h, w, _ = comp.shape out = np.zeros((t, h + FTR, w, 3), np.uint8) out[:, :h] = comp has_crit = 0 <= crit_idx < t tl0, tl1, y = int(w * 0.45), w - 14, FTR // 2 for i in range(t): bar = np.full((FTR, w, 3), 22, np.uint8) is_crit = has_crit and i == crit_idx text = f"frame {i:02d}/{t - 1}" + (" CRITICAL FRAME" if is_crit else "") cv2.putText(bar, text, (8, FTR - 8), FONT, 0.5, CRIT_COLOR if is_crit else (200, 200, 200), 1, cv2.LINE_AA) cv2.line(bar, (tl0, y), (tl1, y), (70, 70, 70), 2, cv2.LINE_AA) if has_crit: cx = int(tl0 + (tl1 - tl0) * crit_idx / max(t - 1, 1)) cv2.line(bar, (cx, y - 7), (cx, y + 7), CRIT_COLOR, 2, cv2.LINE_AA) px = int(tl0 + (tl1 - tl0) * i / max(t - 1, 1)) cv2.circle(bar, (px, y), 4, (255, 255, 255), -1, cv2.LINE_AA) out[i, h:] = bar if is_crit: cv2.rectangle(out[i], (0, 0), (w - 1, h + FTR - 1), CRIT_COLOR, 3) # yuv420p needs even dimensions ph, pw = out.shape[1] % 2, out.shape[2] % 2 if ph or pw: out = np.pad(out, ((0, 0), (0, ph), (0, pw), (0, 0))) return out def write_mp4(frames: np.ndarray, path: Path, fps: float) -> str: _, h, w, _ = frames.shape writer = imageio_ffmpeg.write_frames( str(path), size=(w, h), fps=fps, codec="libx264", pix_fmt_in="rgb24", pix_fmt_out="yuv420p", macro_block_size=1, output_params=["-crf", "16", "-preset", "veryfast", "-movflags", "+faststart"], ) writer.send(None) for f in frames: writer.send(np.ascontiguousarray(f).tobytes()) writer.close() return str(path) def render_video(session: Session, episode_id: str, view: int, show_diff: bool, fps: float, crit: int) -> str: path = session.render / f"{episode_id}_{view}_{int(show_diff)}_{fps:g}.mp4" if not path.exists(): write_mp4(session.composed(episode_id, view, show_diff, crit), path, fps) return str(path) # --------------------------------------------------------------------------- ui glue def metrics_md(row: dict) -> str: def fmt(v, unit=""): return f"{v:.2f}{unit}" if isinstance(v, (int, float)) else "-" out = [ f"**{row['task'].replace('_', ' ')}**", "", f"`{row['episode_id']}`", "", "| | full window | critical frame |", "|---|---|---|", f"| **PSNR** | {fmt(row['full_psnr'], ' dB')} | {fmt(row['crit_psnr'], ' dB')} |", f"| **MSE** | {fmt(row['full_mse'])} | {fmt(row['crit_mse'])} |", "", ] bits = [b for b in (row["edge"], row["arm"]) if b and b != "?"] if row["critical_frame"] is not None: bits.append(f"critical frame {row['critical_frame']}") if bits: out.append(" | ".join(bits)) return "\n".join(out) def critical_index(row: dict) -> int: start, crit = row.get("start_frame"), row.get("critical_frame") if start is None or crit is None: return NO_CRIT return int(crit) - int(start) def on_load(repo_text, sid): """Resolve a repo and load only its index - no window data is pulled here.""" sweep() session = get_session(sid) repo, revision = parse_repo(repo_text) if not repo: return session.sid, "Enter a dataset repo id, e.g. `owner/name`." session.set_repo(repo, revision) try: note = build_index(session) except Exception as exc: session.rows = [] gr.Warning(f"Could not load {repo}: {exc}") return session.sid, f"**{repo}** failed: {exc}" at = f"@{revision}" if revision else "" return session.sid, f"**{repo}{at}** - {len(session.rows)} episodes from {note}" def on_filter(sid, query, sort, current): session = get_session(sid) if not session.rows: return gr.update(choices=[], value=None), [] rows = filter_rows(session, query, sort) choices = [(label_for(r), r["episode_id"]) for r in rows] ids = [r["episode_id"] for r in rows] value = current if current in ids else (ids[0] if ids else None) return ( gr.update(choices=choices, value=value, label=f"Episode ({len(ids)} of {len(session.rows)})"), ids, ) def on_select(sid, episode_id, view, show_diff, fps): session = get_session(sid) if not session.rows or not episode_id: return None, "*No episode selected.*", gr.update() row = row_by_id(session, episode_id) crit = critical_index(row) choices = camera_choices(n_cameras(session.pair(episode_id)[0])) valid = [v for _, v in choices] view = view if view in valid else choices[0][1] video = render_video(session, episode_id, view, show_diff, fps, crit) return video, metrics_md(row), gr.update(choices=choices, value=view) def step(ids, current, delta): if not ids: return None try: i = ids.index(current) except ValueError: i = 0 return ids[(i + delta) % len(ids)] # --------------------------------------------------------------------------- app CSS = """ .window-video video { max-height: 82vh; object-fit: contain; background: #111; } """ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(600, 600)) as demo: sid_state = gr.State("", time_to_live=SESSION_TTL, delete_callback=release) ids_state = gr.State([]) with gr.Row(): with gr.Column(scale=1, min_width=340): gr.Markdown("### Rollout vs. ground truth") with gr.Row(): repo_box = gr.Textbox( label="Dataset repo", placeholder="owner/name (or owner/name@revision, or a dataset URL)", scale=4, ) load_btn = gr.Button("Load", variant="primary", scale=1, min_width=70) status = gr.Markdown("Enter a dataset repo id above and press **Load**.") query = gr.Textbox(label="Search", placeholder="task or episode id") sort = gr.Dropdown(SORTS, value=SORTS[0], label="Sort") picker = gr.Dropdown([], label="Episode", filterable=True) with gr.Row(): prev_btn = gr.Button("Prev") next_btn = gr.Button("Next") view = gr.Radio(camera_choices(3), value=0, label="Camera") show_diff = gr.Checkbox(False, label="Add |GT - generated| panel") fps = gr.Slider(2, 30, value=8, step=1, label="Playback fps") info = gr.Markdown() with gr.Column(scale=4): video = gr.Video( show_label=False, # the panels are captioned in-frame; the chip covers them autoplay=True, loop=True, buttons=["download", "fullscreen"], elem_classes="window-video", ) filter_io = ([sid_state, query, sort, picker], [picker, ids_state]) select_io = ([sid_state, picker, view, show_diff, fps], [video, info, view]) # Load and filter only repopulate the episode list; selecting an episode is what renders. for event in (load_btn.click, repo_box.submit): event(on_load, [repo_box, sid_state], [sid_state, status]).then(on_filter, *filter_io) for control in (query, sort): control.change(on_filter, *filter_io) picker.change(on_select, *select_io) for control in (view, show_diff, fps): control.change(on_select, *select_io) prev_btn.click(lambda i, c: step(i, c, -1), [ids_state, picker], picker) next_btn.click(lambda i, c: step(i, c, +1), [ids_state, picker], picker) if __name__ == "__main__": demo.queue(max_size=16).launch(css=CSS)