robometer/rewardeval_ui
4
1#!/usr/bin/env python32"""3Utility functions for visualization in Robometer (RBM) evaluations.4"""5 6from typing import Optional7import os8import logging9import tempfile10import numpy as np11import matplotlib12matplotlib.use("Agg")13import matplotlib.pyplot as plt14import matplotlib.ticker as ticker15import decord16 17logger = logging.getLogger(__name__)18 19# Colors and layout for progress/success animation (Robometer red)20PROGRESS_COLOR = "#B20000"21SUCCESS_COLOR = "#B20000"22THEME_LIGHT = {"facecolor": "white", "text_color": "black", "spine_color": "#333333"}23 24# Serif font (Palatino) for plots25plt.rcParams["font.family"] = "serif"26plt.rcParams["font.serif"] = ["Palatino", "Palatino Linotype", "DejaVu Serif", "serif"]27plt.rcParams["font.size"] = 1128 29 30def create_combined_progress_success_plot(31 progress_pred: np.ndarray,32 num_frames: int,33 success_binary: Optional[np.ndarray] = None,34 success_probs: Optional[np.ndarray] = None,35 success_labels: Optional[np.ndarray] = None,36 is_discrete_mode: bool = False,37 title: Optional[str] = None,38 loss: Optional[float] = None,39 pearson: Optional[float] = None,40) -> plt.Figure:41 """Create a combined plot with progress, success binary, and success probabilities.42 43 This function creates a unified plot with 1 subplot (progress only) or 3 subplots44 (progress, success binary, success probs), similar to the one used in compile_results.py.45 46 Args:47 progress_pred: Progress predictions array48 num_frames: Number of frames49 success_binary: Optional binary success predictions50 success_probs: Optional success probability predictions51 success_labels: Optional ground truth success labels52 is_discrete_mode: Whether progress is in discrete mode (deprecated, kept for compatibility)53 title: Optional title for the plot (if None, auto-generated from loss/pearson)54 loss: Optional loss value to display in title55 pearson: Optional pearson correlation to display in title56 57 Returns:58 matplotlib Figure object59 """60 # Determine if we should show success plots61 has_success_binary = success_binary is not None and len(success_binary) == len(progress_pred)62 63 if has_success_binary:64 # Three subplots: progress, success (binary), success_probs65 fig, axs = plt.subplots(1, 3, figsize=(18, 3.5))66 ax = axs[0] # Progress subplot67 ax2 = axs[1] # Success subplot (binary)68 ax3 = axs[2] # Success probs subplot69 else:70 # Single subplot: progress only71 fig, ax = plt.subplots(figsize=(7, 3.5))72 ax2 = None73 ax3 = None74 75 # Plot progress76 ax.plot(progress_pred, linewidth=2)77 ax.set_ylabel("Progress")78 79 # Build title80 if title is None:81 title_parts = ["Progress"]82 if loss is not None:83 title_parts.append(f"Loss: {loss:.3f}")84 if pearson is not None:85 title_parts.append(f"Pearson: {pearson:.2f}")86 title = ", ".join(title_parts)87 fig.suptitle(title)88 89 # Set y-limits and ticks (always continuous since discrete is converted before this function)90 ax.set_ylim(0, 1)91 ax.spines["right"].set_visible(False)92 ax.spines["top"].set_visible(False)93 y_ticks = [0, 0.2, 0.4, 0.6, 0.8, 1.0]94 ax.set_yticks(y_ticks)95 96 # Setup success binary subplot97 if ax2 is not None:98 ax2.step(range(len(success_binary)), success_binary, where="post", linewidth=2, label="Predicted", color="blue")99 # Add ground truth success labels as green line if available100 if success_labels is not None and len(success_labels) == len(success_binary):101 ax2.step(102 range(len(success_labels)),103 success_labels,104 where="post",105 linewidth=2,106 label="Ground Truth",107 color="green",108 )109 ax2.set_ylabel("Success (Binary)")110 ax2.set_ylim(-0.05, 1.05)111 ax2.spines["right"].set_visible(False)112 ax2.spines["top"].set_visible(False)113 ax2.set_yticks([0, 1])114 ax2.legend()115 116 # Setup success probs subplot if available117 if ax3 is not None and success_probs is not None:118 ax3.plot(range(len(success_probs)), success_probs, linewidth=2, label="Success Prob", color="purple")119 # Add ground truth success labels as green line if available120 if success_labels is not None and len(success_labels) == len(success_probs):121 ax3.step(122 range(len(success_labels)),123 success_labels,124 where="post",125 linewidth=2,126 label="Ground Truth",127 color="green",128 linestyle="--",129 )130 ax3.set_ylabel("Success Probability")131 ax3.set_ylim(-0.05, 1.05)132 ax3.spines["right"].set_visible(False)133 ax3.spines["top"].set_visible(False)134 ax3.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])135 ax3.legend()136 137 plt.tight_layout()138 return fig139 140 141def extract_frames(video_path: str, fps: float = 1.0, max_frames: int = 64) -> np.ndarray:142 """Extract frames from video file as numpy array (T, H, W, C).143 144 Supports both local file paths and URLs (e.g., HuggingFace Hub URLs).145 Uses the provided ``fps`` to control how densely frames are sampled from146 the underlying video, but caps the total number of frames at ``max_frames``147 to prevent memory issues.148 149 Args:150 video_path: Path to video file or URL151 fps: Frames per second to extract (default: 1.0)152 max_frames: Maximum number of frames to extract (default: 64). This prevents153 memory issues with long videos or high FPS settings.154 155 Returns:156 numpy array of shape (T, H, W, C) containing extracted frames, or None if error157 """158 if video_path is None:159 return None160 161 if isinstance(video_path, tuple):162 video_path = video_path[0]163 164 # Check if it's a URL or local file165 is_url = video_path.startswith(("http://", "https://"))166 is_local_file = os.path.exists(video_path) if not is_url else False167 168 if not is_url and not is_local_file:169 logger.warning(f"Video path does not exist: {video_path}")170 return None171 172 try:173 # decord.VideoReader can handle both local files and URLs174 vr = decord.VideoReader(video_path, num_threads=1)175 total_frames = len(vr)176 177 # Determine native FPS; fall back to a reasonable default if unavailable178 try:179 native_fps = float(vr.get_avg_fps())180 except Exception:181 native_fps = 1.0182 183 # If user-specified fps is invalid or None, default to native fps184 if fps is None or fps <= 0:185 fps = native_fps186 187 # Compute how many frames we want based on desired fps188 # num_frames ≈ total_duration * fps = total_frames * (fps / native_fps)189 if native_fps > 0:190 desired_frames = int(round(total_frames * (fps / native_fps)))191 else:192 desired_frames = total_frames193 194 # Clamp to [1, total_frames]195 desired_frames = max(1, min(desired_frames, total_frames))196 197 # IMPORTANT: Cap at max_frames to prevent memory issues198 # This is critical when fps is high or videos are long199 if desired_frames > max_frames:200 logger.warning(201 f"Requested {desired_frames} frames but capping at {max_frames} "202 f"to prevent memory issues (video has {total_frames} frames at {native_fps:.2f} fps, "203 f"requested extraction at {fps:.2f} fps)"204 )205 desired_frames = max_frames206 207 # Evenly sample indices to match the desired number of frames208 if desired_frames == total_frames:209 frame_indices = list(range(total_frames))210 else:211 frame_indices = np.linspace(0, total_frames - 1, desired_frames, dtype=int).tolist()212 213 frames_array = vr.get_batch(frame_indices).asnumpy() # Shape: (T, H, W, C)214 del vr215 return frames_array216 except Exception as e:217 logger.error(f"Error extracting frames from {video_path}: {e}")218 return None219 220 221def resize_frames_keep_aspect(222 frames: np.ndarray,223 max_edge: int = 480,224) -> np.ndarray:225 """Resize video frames so the longer edge is at most max_edge, preserving aspect ratio.226 Use when creating videos so the image is not stretched. Uses scipy if available.227 """228 if frames is None or frames.size == 0 or frames.ndim != 4:229 return frames230 t, h, w, c = frames.shape231 if h <= 0 or w <= 0:232 return frames233 scale = min(max_edge / max(h, w), 1.0)234 if scale >= 1.0:235 return frames236 new_h = max(1, round(h * scale))237 new_w = max(1, round(w * scale))238 try:239 from scipy.ndimage import zoom240 zoom_factors = (1.0, new_h / h, new_w / w, 1.0)241 out = zoom(frames.astype(np.float64), zoom_factors, order=1)242 return np.clip(out, 0, 255).astype(np.uint8)243 except ImportError:244 return frames245 246 247def _style_progress_ax(ax, theme: dict, ylabel: str = "Progress"):248 """Style a progress or success axis (shared look)."""249 ax.set_facecolor(theme["facecolor"])250 ax.set_ylim(-0.05, 1.05)251 ax.set_xlabel("")252 ax.set_ylabel(ylabel, fontsize=12, fontweight="bold", color=theme["text_color"])253 ax.spines["left"].set_color(theme["spine_color"])254 ax.spines["bottom"].set_color(theme["spine_color"])255 ax.spines["right"].set_visible(False)256 ax.spines["top"].set_visible(False)257 ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=8))258 ax.set_yticks([0, 0.5, 1.0])259 ax.tick_params(axis="both", labelsize=10, colors=theme["text_color"])260 261 262def create_progress_success_gif(263 progress_pred: np.ndarray,264 success_data: Optional[np.ndarray] = None,265 video_frames: Optional[np.ndarray] = None,266 output_path: Optional[str] = None,267 title: Optional[str] = None,268 duration_sec: float = 5.0,269 theme: Optional[dict] = None,270) -> Optional[str]:271 """Create an animated MP4: progress and success curves growing frame-by-frame (optional video on left).272 273 Uses light theme by default for web UI. Output is always 5 seconds (duration_sec); fps is274 computed as num_frames / duration_sec. Saves to output_path as .mp4. Returns path if saved, None on error.275 """276 from matplotlib.animation import FuncAnimation277 278 theme = theme or THEME_LIGHT279 progress_pred = np.atleast_1d(progress_pred).astype(float)280 num_frames = len(progress_pred)281 if num_frames == 0:282 return None283 284 # FPS so the full animation runs for duration_sec (e.g. 5 seconds)285 fps = max(1, round(num_frames / duration_sec))286 287 success_padded = None288 if success_data is not None and np.size(success_data) > 0:289 s = np.atleast_1d(success_data).astype(float)290 if len(s) < num_frames:291 s = np.pad(s, (0, num_frames - len(s)), mode="edge")292 success_padded = s293 294 has_video = (295 video_frames is not None296 and getattr(video_frames, "shape", (0,))[0] >= num_frames297 )298 if has_video and video_frames.shape[0] > num_frames:299 video_frames = video_frames[:num_frames]300 elif has_video and video_frames.shape[0] < num_frames:301 pad = np.repeat(video_frames[-1:], num_frames - video_frames.shape[0], axis=0)302 video_frames = np.concatenate([video_frames, pad], axis=0)303 if has_video:304 video_frames = resize_frames_keep_aspect(video_frames, max_edge=480)305 306 n_panels = 2 if success_padded is not None else 1307 width_per_panel = 5.5308 figsize = (width_per_panel * n_panels, 3.2) if not has_video else (2 + width_per_panel * n_panels, 3.2)309 310 if has_video:311 from matplotlib.gridspec import GridSpec312 fig = plt.figure(facecolor=theme["facecolor"], figsize=figsize)313 # Give plots more room: smaller video column, more wspace so video doesn't cover Progress314 gs = GridSpec(1, 2, figure=fig, width_ratios=[0.85, n_panels], wspace=0.4)315 ax_video = fig.add_subplot(gs[0])316 ax_video.set_facecolor(theme["facecolor"])317 ax_video.axis("off")318 # Preserve aspect ratio so the video is not flattened319 vid_im = ax_video.imshow(320 np.clip(video_frames[0], 0, 255).astype(np.uint8)321 if video_frames[0].ndim >= 3322 else video_frames[0],323 cmap="gray" if video_frames[0].ndim == 2 else None,324 aspect="equal",325 )326 from matplotlib.gridspec import GridSpecFromSubplotSpec327 gs_right = GridSpecFromSubplotSpec(1, n_panels, subplot_spec=gs[1], wspace=0.3)328 axes = [fig.add_subplot(gs_right[0, j]) for j in range(n_panels)]329 else:330 fig, axes = plt.subplots(331 1, n_panels, figsize=figsize, facecolor=theme["facecolor"]332 )333 axes = np.atleast_1d(axes)334 vid_im = None335 336 lines = []337 head_dots = []338 for i in range(n_panels):339 ax = axes[i]340 if i == 1 and success_padded is not None:341 _style_progress_ax(ax, theme, ylabel="Success")342 ax.set_xlim(-0.5, num_frames)343 line, = ax.plot([], [], lw=2.5, color=SUCCESS_COLOR, drawstyle="steps-post")344 lines.append(line)345 head_dots.append(None)346 else:347 _style_progress_ax(ax, theme, ylabel="Progress")348 ax.set_xlim(-0.5, num_frames)349 line, = ax.plot([], [], lw=2.5, color=PROGRESS_COLOR, drawstyle="steps-post")350 head_dot = ax.scatter(351 [], [], color=PROGRESS_COLOR, s=36, zorder=5,352 edgecolors=PROGRESS_COLOR, facecolors="none",353 )354 lines.append(line)355 head_dots.append(head_dot)356 357 if title and str(title).strip():358 # Place title inside figure top margin (rect keeps axes below 0.88)359 fig.suptitle(360 str(title).strip(),361 fontsize=12,362 fontweight="bold",363 color=theme["text_color"],364 y=0.94,365 )366 367 def update(frame):368 out = []369 if vid_im is not None and has_video:370 idx = min(int(frame), video_frames.shape[0] - 1)371 f = np.clip(video_frames[idx], 0, 255).astype(np.uint8)372 if f.ndim == 2:373 vid_im.set_cmap("gray")374 vid_im.set_array(f)375 out.append(vid_im)376 for i in range(n_panels):377 if i == 1 and success_padded is not None:378 x = np.arange(int(frame) + 1)379 y = success_padded[: int(frame) + 1]380 if len(x) > 0 and len(y) > 0:381 lines[i].set_data(x, y)382 else:383 x = np.arange(int(frame) + 1)384 y = progress_pred[: int(frame) + 1]385 if len(x) > 0 and len(y) > 0:386 lines[i].set_data(x, y)387 if head_dots[i] is not None:388 head_dots[i].set_offsets([[frame, progress_pred[int(frame)]]])389 out.append(lines[i])390 if head_dots[i] is not None:391 out.append(head_dots[i])392 return out393 394 # Leave extra top space so suptitle (task text) is not cut off; minimal horizontal pad for tight video395 plt.tight_layout(rect=[0.01, 0, 0.99, 0.88], pad=0.3)396 ani = FuncAnimation(397 fig, update, frames=num_frames, interval=1000 / fps, blit=True398 )399 400 if not output_path:401 fd, output_path = tempfile.mkstemp(suffix=".mp4")402 os.close(fd)403 # Normalize to .mp4404 if output_path.endswith(".gif"):405 output_path = output_path[:-4] + ".mp4"406 if not output_path.lower().endswith(".mp4"):407 output_path = output_path + ".mp4"408 out_dir = os.path.dirname(output_path)409 if out_dir:410 os.makedirs(out_dir, exist_ok=True)411 412 savefig_kwargs = {413 "facecolor": theme["facecolor"],414 "edgecolor": "none",415 "bbox_inches": "tight",416 "pad_inches": 0.12,417 }418 try:419 ani.save(420 output_path,421 writer="ffmpeg",422 fps=fps,423 dpi=120,424 savefig_kwargs=savefig_kwargs,425 )426 except Exception as e:427 logger.warning(f"Could not save MP4 (ffmpeg?): {e}")428 output_path = None429 finally:430 plt.close(fig)431 432 return output_path433