update

app.py

@@ -27,7 +27,11 @@ from typing import Any, List, Optional, Tuple
 
 from dataset_types import Trajectory, ProgressSample, PreferenceSample
 from eval_utils import build_payload, post_batch_npy
-from eval_viz_utils import create_combined_progress_success_plot, extract_frames
+from eval_viz_utils import (
+    create_combined_progress_success_plot,
+    create_progress_success_gif,
+    extract_frames,
+)
 from datasets import load_dataset as load_dataset_hf, get_dataset_config_names
 
 logger = logging.getLogger(__name__)
@@ -421,22 +425,24 @@ def process_single_video(
     server_url: str = "",
     fps: float = 1.0,
     use_frame_steps: bool = False,
-) -> Tuple[Optional[str], Optional[str]]:
-    """Process single video for progress and success predictions using eval server."""
+) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+    """Process single video for progress and success predictions using eval server.
+    Returns (static_plot_path, video_path, info_text). video_path is the 5 sec MP4 animation; may be None if creation fails.
+    """
     # Get server URL from state if not provided
     if not server_url:
         server_url = _server_state.get("server_url")
 
     if not server_url:
-        return None, "Please select a model from the dropdown above and ensure it's connected."
+        return None, None, "Please select a model from the dropdown above and ensure it's connected."
 
     if video_path is None:
-        return None, "Please provide a video."
+        return None, None, "Please provide a video."
 
     try:
         frames_array = extract_frames(video_path, fps=fps)
         if frames_array is None or frames_array.size == 0:
-            return None, "Could not extract frames from video."
+            return None, None, "Could not extract frames from video."
 
         # Convert frames to (T, H, W, C) numpy array with uint8 values
         if frames_array.dtype != np.uint8:
@@ -520,11 +526,24 @@ def process_single_video(
         if success_array is not None and len(success_array) > 0:
             info_text += f"**Final success probability:** {success_array[-1]:.3f}\n"
 
-        # Return combined plot (which includes success if available)
-        return progress_plot, info_text
+        # Animated MP4: progress + success curves (5 sec clip) with optional video panel
+        video_path = None
+        if len(progress_array) > 0:
+            mp4_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
+            mp4_file.close()
+            video_path = create_progress_success_gif(
+                progress_pred=progress_array,
+                success_data=success_binary if success_binary is not None else success_array,
+                video_frames=frames_array,
+                output_path=mp4_file.name,
+                title=task_text,
+                duration_sec=5.0,
+            )
+
+        return progress_plot, video_path, info_text
 
     except Exception as e:
-        return None, f"Error processing video: {str(e)}"
+        return None, None, f"Error processing video: {str(e)}"
 
 
 def process_two_videos(
@@ -781,6 +800,7 @@ with demo:
 
                 with gr.Column():
                     progress_plot = gr.Image(label="Progress & Success Prediction", height=320)
+                    progress_video = gr.Video(label="Animated Progress & Success (5 sec MP4)", height=320)
                     info_output = gr.Markdown("")
                     gr.Markdown("---")
                     gr.Markdown("**Examples**")
@@ -1004,7 +1024,7 @@ with demo:
                     fps_input_single,
                     use_frame_steps_single,
                 ],
-                outputs=[progress_plot, info_output],
+                outputs=[progress_plot, progress_video, info_output],
                 api_name="process_single_video",
             )
 

eval_viz_utils.py

@@ -8,11 +8,24 @@ import os
 import logging
 import tempfile
 import numpy as np
+import matplotlib
+matplotlib.use("Agg")
 import matplotlib.pyplot as plt
+import matplotlib.ticker as ticker
 import decord
 
 logger = logging.getLogger(__name__)
 
+# Colors and layout for progress/success animation (Robometer red)
+PROGRESS_COLOR = "#B20000"
+SUCCESS_COLOR = "#B20000"
+THEME_LIGHT = {"facecolor": "white", "text_color": "black", "spine_color": "#333333"}
+
+# Serif font (Palatino) for plots
+plt.rcParams["font.family"] = "serif"
+plt.rcParams["font.serif"] = ["Palatino", "Palatino Linotype", "DejaVu Serif", "serif"]
+plt.rcParams["font.size"] = 11
+
 
 def create_combined_progress_success_plot(
     progress_pred: np.ndarray,
@@ -49,13 +62,13 @@ def create_combined_progress_success_plot(
 
     if has_success_binary:
         # Three subplots: progress, success (binary), success_probs
-        fig, axs = plt.subplots(1, 3, figsize=(15, 3.5))
+        fig, axs = plt.subplots(1, 3, figsize=(18, 3.5))
         ax = axs[0]  # Progress subplot
         ax2 = axs[1]  # Success subplot (binary)
         ax3 = axs[2]  # Success probs subplot
     else:
         # Single subplot: progress only
-        fig, ax = plt.subplots(figsize=(6, 3.5))
+        fig, ax = plt.subplots(figsize=(7, 3.5))
         ax2 = None
         ax3 = None
 
@@ -203,3 +216,217 @@ def extract_frames(video_path: str, fps: float = 1.0, max_frames: int = 64) -> n
     except Exception as e:
         logger.error(f"Error extracting frames from {video_path}: {e}")
         return None
+
+
+def resize_frames_keep_aspect(
+    frames: np.ndarray,
+    max_edge: int = 480,
+) -> np.ndarray:
+    """Resize video frames so the longer edge is at most max_edge, preserving aspect ratio.
+    Use when creating videos so the image is not stretched. Uses scipy if available.
+    """
+    if frames is None or frames.size == 0 or frames.ndim != 4:
+        return frames
+    t, h, w, c = frames.shape
+    if h <= 0 or w <= 0:
+        return frames
+    scale = min(max_edge / max(h, w), 1.0)
+    if scale >= 1.0:
+        return frames
+    new_h = max(1, round(h * scale))
+    new_w = max(1, round(w * scale))
+    try:
+        from scipy.ndimage import zoom
+        zoom_factors = (1.0, new_h / h, new_w / w, 1.0)
+        out = zoom(frames.astype(np.float64), zoom_factors, order=1)
+        return np.clip(out, 0, 255).astype(np.uint8)
+    except ImportError:
+        return frames
+
+
+def _style_progress_ax(ax, theme: dict, ylabel: str = "Progress"):
+    """Style a progress or success axis (shared look)."""
+    ax.set_facecolor(theme["facecolor"])
+    ax.set_ylim(-0.05, 1.05)
+    ax.set_xlabel("")
+    ax.set_ylabel(ylabel, fontsize=12, fontweight="bold", color=theme["text_color"])
+    ax.spines["left"].set_color(theme["spine_color"])
+    ax.spines["bottom"].set_color(theme["spine_color"])
+    ax.spines["right"].set_visible(False)
+    ax.spines["top"].set_visible(False)
+    ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=8))
+    ax.set_yticks([0, 0.5, 1.0])
+    ax.tick_params(axis="both", labelsize=10, colors=theme["text_color"])
+
+
+def create_progress_success_gif(
+    progress_pred: np.ndarray,
+    success_data: Optional[np.ndarray] = None,
+    video_frames: Optional[np.ndarray] = None,
+    output_path: Optional[str] = None,
+    title: Optional[str] = None,
+    duration_sec: float = 5.0,
+    theme: Optional[dict] = None,
+) -> Optional[str]:
+    """Create an animated MP4: progress and success curves growing frame-by-frame (optional video on left).
+
+    Uses light theme by default for web UI. Output is always 5 seconds (duration_sec); fps is
+    computed as num_frames / duration_sec. Saves to output_path as .mp4. Returns path if saved, None on error.
+    """
+    from matplotlib.animation import FuncAnimation
+
+    theme = theme or THEME_LIGHT
+    progress_pred = np.atleast_1d(progress_pred).astype(float)
+    num_frames = len(progress_pred)
+    if num_frames == 0:
+        return None
+
+    # FPS so the full animation runs for duration_sec (e.g. 5 seconds)
+    fps = max(1, round(num_frames / duration_sec))
+
+    success_padded = None
+    if success_data is not None and np.size(success_data) > 0:
+        s = np.atleast_1d(success_data).astype(float)
+        if len(s) < num_frames:
+            s = np.pad(s, (0, num_frames - len(s)), mode="edge")
+        success_padded = s
+
+    has_video = (
+        video_frames is not None
+        and getattr(video_frames, "shape", (0,))[0] >= num_frames
+    )
+    if has_video and video_frames.shape[0] > num_frames:
+        video_frames = video_frames[:num_frames]
+    elif has_video and video_frames.shape[0] < num_frames:
+        pad = np.repeat(video_frames[-1:], num_frames - video_frames.shape[0], axis=0)
+        video_frames = np.concatenate([video_frames, pad], axis=0)
+    if has_video:
+        video_frames = resize_frames_keep_aspect(video_frames, max_edge=480)
+
+    n_panels = 2 if success_padded is not None else 1
+    width_per_panel = 5.5
+    figsize = (width_per_panel * n_panels, 3.2) if not has_video else (2 + width_per_panel * n_panels, 3.2)
+
+    if has_video:
+        from matplotlib.gridspec import GridSpec
+        fig = plt.figure(facecolor=theme["facecolor"], figsize=figsize)
+        # Give plots more room: smaller video column, more wspace so video doesn't cover Progress
+        gs = GridSpec(1, 2, figure=fig, width_ratios=[0.85, n_panels], wspace=0.4)
+        ax_video = fig.add_subplot(gs[0])
+        ax_video.set_facecolor(theme["facecolor"])
+        ax_video.axis("off")
+        # Preserve aspect ratio so the video is not flattened
+        vid_im = ax_video.imshow(
+            np.clip(video_frames[0], 0, 255).astype(np.uint8)
+            if video_frames[0].ndim >= 3
+            else video_frames[0],
+            cmap="gray" if video_frames[0].ndim == 2 else None,
+            aspect="equal",
+        )
+        from matplotlib.gridspec import GridSpecFromSubplotSpec
+        gs_right = GridSpecFromSubplotSpec(1, n_panels, subplot_spec=gs[1], wspace=0.3)
+        axes = [fig.add_subplot(gs_right[0, j]) for j in range(n_panels)]
+    else:
+        fig, axes = plt.subplots(
+            1, n_panels, figsize=figsize, facecolor=theme["facecolor"]
+        )
+        axes = np.atleast_1d(axes)
+        vid_im = None
+
+    lines = []
+    head_dots = []
+    for i in range(n_panels):
+        ax = axes[i]
+        if i == 1 and success_padded is not None:
+            _style_progress_ax(ax, theme, ylabel="Success")
+            ax.set_xlim(-0.5, num_frames)
+            line, = ax.plot([], [], lw=2.5, color=SUCCESS_COLOR, drawstyle="steps-post")
+            lines.append(line)
+            head_dots.append(None)
+        else:
+            _style_progress_ax(ax, theme, ylabel="Progress")
+            ax.set_xlim(-0.5, num_frames)
+            line, = ax.plot([], [], lw=2.5, color=PROGRESS_COLOR, drawstyle="steps-post")
+            head_dot = ax.scatter(
+                [], [], color=PROGRESS_COLOR, s=36, zorder=5,
+                edgecolors=PROGRESS_COLOR, facecolors="none",
+            )
+            lines.append(line)
+            head_dots.append(head_dot)
+
+    if title and str(title).strip():
+        # Place title inside figure top margin (rect keeps axes below 0.88)
+        fig.suptitle(
+            str(title).strip(),
+            fontsize=12,
+            fontweight="bold",
+            color=theme["text_color"],
+            y=0.94,
+        )
+
+    def update(frame):
+        out = []
+        if vid_im is not None and has_video:
+            idx = min(int(frame), video_frames.shape[0] - 1)
+            f = np.clip(video_frames[idx], 0, 255).astype(np.uint8)
+            if f.ndim == 2:
+                vid_im.set_cmap("gray")
+            vid_im.set_array(f)
+            out.append(vid_im)
+        for i in range(n_panels):
+            if i == 1 and success_padded is not None:
+                x = np.arange(int(frame) + 1)
+                y = success_padded[: int(frame) + 1]
+                if len(x) > 0 and len(y) > 0:
+                    lines[i].set_data(x, y)
+            else:
+                x = np.arange(int(frame) + 1)
+                y = progress_pred[: int(frame) + 1]
+                if len(x) > 0 and len(y) > 0:
+                    lines[i].set_data(x, y)
+                    if head_dots[i] is not None:
+                        head_dots[i].set_offsets([[frame, progress_pred[int(frame)]]])
+            out.append(lines[i])
+            if head_dots[i] is not None:
+                out.append(head_dots[i])
+        return out
+
+    # Leave extra top space so suptitle (task text) is not cut off; minimal horizontal pad for tight video
+    plt.tight_layout(rect=[0.01, 0, 0.99, 0.88], pad=0.3)
+    ani = FuncAnimation(
+        fig, update, frames=num_frames, interval=1000 / fps, blit=True
+    )
+
+    if not output_path:
+        fd, output_path = tempfile.mkstemp(suffix=".mp4")
+        os.close(fd)
+    # Normalize to .mp4
+    if output_path.endswith(".gif"):
+        output_path = output_path[:-4] + ".mp4"
+    if not output_path.lower().endswith(".mp4"):
+        output_path = output_path + ".mp4"
+    out_dir = os.path.dirname(output_path)
+    if out_dir:
+        os.makedirs(out_dir, exist_ok=True)
+
+    savefig_kwargs = {
+        "facecolor": theme["facecolor"],
+        "edgecolor": "none",
+        "bbox_inches": "tight",
+        "pad_inches": 0.12,
+    }
+    try:
+        ani.save(
+            output_path,
+            writer="ffmpeg",
+            fps=fps,
+            dpi=120,
+            savefig_kwargs=savefig_kwargs,
+        )
+    except Exception as e:
+        logger.warning(f"Could not save MP4 (ffmpeg?): {e}")
+        output_path = None
+    finally:
+        plt.close(fig)
+
+    return output_path