Fall back to the composed side-by-side video

Rerun could not be made to work: the gradio_rerun viewer intermittently
stops mid-update, reproducible in a 40-line app with synthetic data, and
pinning the reference Space's gradio 6.5.1 broke the page outright.

Keeps the simple single-view layout -- episode list, search, sort, PSNR/MSE --
and puts back the composed mp4, which frame-locks ground truth against the
generated rollout in one player. Camera selector, optional difference panel
and playback fps are retained; drops rerun-sdk and gradio_rerun.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

app.py

@@ -1,5 +1,5 @@
 """
-Rerun viewer for forward-dynamics rollouts vs. ground truth.
+Side-by-side viewer for forward-dynamics rollouts vs. ground truth.
 
 Point it at any HF dataset repo laid out as:
 
@@ -8,15 +8,16 @@ Point it at any HF dataset repo laid out as:
     windows/<episode_id>/gt.mp4         # ground-truth window
     windows/<episode_id>/generated.mp4  # model rollout for the same slice
 
-Pick an episode and its two videos stream into an embedded Rerun viewer, side
-by side on a shared frame timeline.
+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. Recordings
-go straight to the viewer as a binary stream, so no .rrd files are written.
+MAX_WINDOWS_ON_DISK windows, and is deleted when the session ends.
 
-The app starts empty and assumes no dataset: type a repo id and press Load.
 Private repos need an HF_TOKEN secret with read access.
 """
 
@@ -35,10 +36,8 @@ from pathlib import Path
 
 import cv2
 import gradio as gr
+import imageio_ffmpeg
 import numpy as np
-import rerun as rr
-import rerun.blueprint as rrb
-from gradio_rerun import Rerun
 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
@@ -48,9 +47,22 @@ 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))
-JPEG_QUALITY = 90
+
+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)",
@@ -90,6 +102,7 @@ class Session:
         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:
@@ -123,6 +136,8 @@ class Session:
             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)
 
@@ -143,6 +158,18 @@ class Session:
             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()
@@ -151,6 +178,7 @@ class Session:
     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)
@@ -158,6 +186,7 @@ class Session:
 
     def close(self) -> None:
         self._pairs.clear()
+        self._composed.clear()
         shutil.rmtree(self.root, ignore_errors=True)
 
 
@@ -339,6 +368,143 @@ def read_video(path: str) -> np.ndarray:
     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
 
 
@@ -368,7 +534,7 @@ def metrics_md(row: dict) -> str:
 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 -1
+        return NO_CRIT
     return int(crit) - int(start)
 
 
@@ -404,50 +570,19 @@ def on_filter(sid, query, sort, current):
     )
 
 
-def on_metrics(sid, episode_id):
+def on_select(sid, episode_id, view, show_diff, fps):
     session = get_session(sid)
     if not session.rows or not episode_id:
-        return "*No episode selected.*"
-    return metrics_md(row_by_id(session, 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]
 
-def render_episode(sid, episode_id):
-    """Write the episode to an .rrd: both videos side by side on one frame timeline.
-
-    Handing the viewer a finished file rather than a live binary stream: streaming
-    races the viewer's wasm boot, and chunks that land before it is ready are dropped
-    against a stopped instance.
-    """
-    session = get_session(sid)
-    if not session.rows or not episode_id:
-        # Never hand the viewer a null: that resets and stops the wasm instance, and the
-        # next real value then tries to open a channel on a dead viewer.
-        return gr.skip()
-
-    path = session.render / f"{episode_id}.rrd"
-    if path.exists():
-        return str(path)
-
-    gt, gen = session.pair(episode_id)
-    crit = critical_index(row_by_id(session, episode_id))
-    rec = rr.RecordingStream("wmviz", recording_id=episode_id)
-    for i in range(len(gt)):
-        rec.set_time("frame", sequence=i)
-        rec.log("ground_truth", rr.Image(gt[i]).compress(jpeg_quality=JPEG_QUALITY))
-        rec.log("generated", rr.Image(gen[i]).compress(jpeg_quality=JPEG_QUALITY))
-        if i == crit:
-            rec.log("critical_frame", rr.TextLog("critical frame"))
-    rec.save(
-        str(path),
-        default_blueprint=rrb.Blueprint(
-            rrb.Horizontal(
-                rrb.Spatial2DView(origin="ground_truth", name="Ground truth"),
-                rrb.Spatial2DView(origin="generated", name="Generated"),
-            ),
-            collapse_panels=True,
-        ),
-    )
-    return str(path)
+    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):
@@ -462,6 +597,10 @@ def step(ids, current, delta):
 
 # --------------------------------------------------------------------------- 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([])
@@ -483,44 +622,36 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(
             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):
-            # Config matched to rerun/gradio-rerun-viewer: a CSS-string height leaves the
-            # container at zero height while the viewer re-initialises, and it never recovers.
-            viewer = Rerun(
-                streaming=True,
-                height=820,
-                panel_states={"time": "expanded", "blueprint": "hidden", "selection": "hidden"},
+            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])
 
-    # Loading and filtering only repopulate the episode list. Selecting an episode is the
-    # single trigger that streams into the viewer -- if these chains also streamed, setting
-    # the dropdown value here would fire picker.change too and run two generators
-    # concurrently into one viewer, which crashes it.
+    # 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_metrics, [sid_state, picker], info)
-    # trigger_mode="once" is load-bearing: .change() defaults to "always_last", which lets a
-    # second run queue up behind the pending one. Two generators streaming into one viewer
-    # stops the first viewer mid-write, which is what left it stuck on "Loading Rerun...".
-    picker.change(
-        render_episode,
-        [sid_state, picker],
-        viewer,
-        concurrency_limit=1,
-        trigger_mode="once",
-    )
+    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__":
-    # ssr_mode must stay off: server-side rendering re-mounts the viewer canvas on
-    # hydration, which kills the wasm instance mid-stream.
-    demo.queue(max_size=16).launch(ssr_mode=False)
+    demo.queue(max_size=16).launch(css=CSS)

requirements.txt

@@ -1,5 +1,5 @@
 gradio==6.22.0
-gradio_rerun==0.35.0
 huggingface_hub>=0.28
 numpy
 opencv-python-headless
+imageio-ffmpeg