# Created by BVH, Mar - Aug 2025. # Helper methods for visualizations. import os import re import time import cv2 import cv2 import imageio import lovely_tensors import matplotlib import matplotlib.pyplot as plt import numpy as np import skimage.metrics import sklearn.decomposition import torch from einops import rearrange from lovely_numpy import lo from rich import print from imaginaire.utils import log lovely_tensors.monkey_patch() np.set_printoptions(precision=3, suppress=True) torch.set_printoptions(precision=3, sci_mode=False, threshold=1000) from custom.vae.action_codec import actions_to_absolute, unnormalize_actions CAMS_COMB_WIDTH = 32 def tensor_to_thwc_uint8(video): ''' Convert a (C, T, H, W) float tensor in [-1, 1] into a (T, H, W, C) uint8 numpy array in [0, 255]. Idempotent if the input is already a (T, H, W, C) uint8 numpy array (returned as-is). Mirrors the same scale/clip/rearrange flex_save_video does internally, but exposed as a small helper so callers can do an explicit tensor->numpy transition without going through MP4 encoding. ''' if isinstance(video, np.ndarray) and video.dtype == np.uint8 \ and video.ndim == 4 and video.shape[-1] == 3: return video if isinstance(video, torch.Tensor): if 'bfloat' in str(video.dtype).lower(): video = video.type(torch.float32) video = video.detach().cpu().numpy() assert video.ndim == 4, f'Expected 4D input, got {video.shape}' # Assume (C, T, H, W) float layout for tensor-origin input. video = rearrange(video, 'c t h w -> t h w c') if video.dtype.kind == 'f': if video.min() < 0.0: video = (video + 1.0) / 2.0 video = video.clip(0.0, 1.0) if np.isnan(video).any() or np.isinf(video).any(): print(f'[orange3]tensor_to_thwc_uint8(): Warning: NaN or Inf in video?') video = np.nan_to_num(video, nan=0.0, posinf=0.0, neginf=0.0) video = (video * 255.0).astype(np.uint8) else: assert video.dtype.kind in 'ui' assert 0 <= video.min() and video.max() <= 255 video = video.astype(np.uint8) return video def flex_save_video(x, n=2, p=None, d=0, c=2, a=0.0, u=1, m='magma', f=10, h=3, q=8, b=8, r=2, e=1, o=1, v=2): ''' Robust save video method. :param x (array / tensor): (*, C, T, H, W) or (*, T, H, W, C). :param n (bool): Normalize from [-1, 1] to [0, 1]. :param p (str): Path or file name prefix to save the video to or with. :param d (int): Channel dimension. :param c (int): Apply PCA first. :param a (float): Aspect ratio to center crop to. :param u (int): Upscale factor to avoid blurry viewing. :param m (int): If C = 1, apply colormap. :param f (int): Frames per second. :param h (int): Hold last frame for this many extra frames for clearer viewing. :param q (int): Exported video quality. :param b (int): Macro block size. :param r (int): Resize to multiple of macro block size. 0 = off, 1 = crop (round down), 2 = zero-pad (round up, preserves content). :param e (bool): Export video to disk, otherwise just return the visualization. :param o (bool): Overwrite existing mp4 video if it exists. :param v (int): Verbosity. :return (str, array): Path to exported visualization and (T, H, W, C) array thereof. ''' start_time = time.time() # Convert to numpy array with correct shape. if isinstance(x, list): x = np.stack(x) elif isinstance(x, torch.Tensor): if 'bfloat' in str(x.dtype).lower(): x = x.type(torch.float32) x = x.detach().cpu().numpy() while x.ndim < 4: if v >= 2: print(f'[gray]flex_save_video(): Expanding video from {x.shape} to {x[None].shape}') x = x[None] # e.g. (H, W) -> (C, T, H, W) with C=1, T=1. while x.ndim > 4: if v >= 2: print(f'[gray]flex_save_video(): Collapsing video from {x.shape} to {x[0].shape}') x = x[0] # e.g. (B, C, T, H, W) -> (C, T, H, W) assuming B=1. if d in [0, -4]: x = rearrange(x, 'c t h w -> t h w c') elif d in [1, -3]: x = rearrange(x, 't c h w -> t h w c') elif d in [2, -2]: x = rearrange(x, 't h c w -> t h w c') elif d in [3, -1]: pass else: raise ValueError(f'Unexpected channel dimension: {d}') (T, H, W, C) = x.shape # Perform PCA if requested. if c == 1 or (c == 2 and C > 3): if v >= 2: print(f'[gray]flex_save_video(): Applying PCA from {C} to 3 channels') if x.var() != 0.0: old1_x = x x = easy_pca(x, k=3, normalize=[0.0, 1.0]) else: # Treat all constant array differently to avoid warnings. x = np.zeros_like(x[..., 0:3]) (T, H, W, C) = x.shape # Convert grayscale to RGB if needed. if C == 1: if m is None: x = np.repeat(x, 3, axis=-1) else: # Apply colormap. x_min = x.min() x_max = x.max() x = (x - x_min) / (x_max - x_min) x = np.stack([matplotlib.colormaps[m](x[i, :, :, 0]) for i in range(T)])[..., 0:3] (T, H, W, C) = x.shape assert C == 3, f'Expected 3 channels (RGB), got {C}' # Perform last frame freezing if requested. if h > 0: old2_x = x x = np.concatenate([x, x[-1:].repeat(h, axis=0)], axis=0) (T, H, W, C) = x.shape # Handle data type and normalization. if x.dtype.kind == 'f': if n == 1 or (n == 2 and x.min() < 0.0): if v >= 2: print(f'[gray]Normalizing from [-1, 1] to [0, 1]') x = (x + 1.0) / 2.0 x = x.clip(0.0, 1.0) if np.isnan(x).any() or np.isinf(x).any(): print(f'[orange3]flex_save_video(): Warning: NaN or Inf in x?') x[np.isnan(x)] = 0.0 x[np.isinf(x)] = 0.0 x = (x * 255.0).astype(np.uint8) else: assert x.dtype.kind in 'ui' assert 0 <= x.min() and x.max() <= 255 x = x.astype(np.uint8) # Handle aspect ratio cropping if requested. if 0.01 < a < 100.0: a_old = W / H if a > a_old + 3e-3: # Go wider -- crop vertically. H_new = int(W / a) if v >= 2: print(f'[gray]flex_save_video(): Cropping vertically from {H} to {H_new}') offset = (H - H_new) // 2 x = x[:, offset:offset + H_new, :, :] elif a < a_old - 3e-3: # Go narrower -- crop horizontally. W_new = int(H * a) if v >= 2: print(f'[gray]flex_save_video(): Cropping horizontally from {W} to {W_new}') offset = (W - W_new) // 2 x = x[:, :, offset:offset + W_new, :] (T, H, W, C) = x.shape # Perform upscaling if requested; Tile pixels by broadcasting for efficiency. if u > 1: x = x.reshape(T, H, 1, W, 1, C) x = np.broadcast_to(x, (T, H, u, W, u, C)) x = x.reshape(T, H * u, W * u, C) (T, H, W, C) = x.shape # Handle macro block size alignment if requested. if r >= 1 and b > 1: if r == 1: # Crop (round down) -- may discard content at right/bottom edges. H_new = int(H / b) * b W_new = int(W / b) * b x = x[:, 0:H_new, 0:W_new, :] elif r == 2: # Zero-pad (round up) -- preserves all content. H_new = ((H + b - 1) // b) * b W_new = ((W + b - 1) // b) * b if H_new != H or W_new != W: x = np.pad(x, ((0, 0), (0, H_new - H), (0, W_new - W), (0, 0)), mode='constant', constant_values=0) (T, H, W, C) = x.shape # Save to disk only if requested. if e >= 1: # Remove existing file if requested. if o >= 1 and '.mp4' in p and os.path.exists(p) and os.path.isfile(p): if v >= 1: print(f'[orange3]flex_save_video(): Warning: {p} already exists, overwriting...') os.remove(p) # Generate output file path if not clearly provided. else: i = 0 fn = p while p is None or '.mp4' not in p or os.path.exists(p): if fn is not None and '.mp4' not in fn: if '/' not in fn: p = f'out_dbg/{fn}_{i:03d}.mp4' else: p = f'{fn}_{i:03d}.mp4' else: p = f'out_dbg/v_{i:03d}.mp4' i += 1 if v >= 1: print(f'[yellow]flex_save_video(): Saving video to: {p}...') os.makedirs(os.path.dirname(p), exist_ok=True) imageio.mimwrite(p, x, format='ffmpeg', fps=float(f), macro_block_size=b, quality=max(min(int(q), 10), 2)) end_time = time.time() if v >= 2: print(f'[magenta]flex_save_video(): Took {end_time - start_time:.3f} seconds for: {p}') return (p, x) def easy_pca(array, k=3, normalize=[0.0, 1.0]): ''' :param array (*, n): Array to perform PCA on. :param k (int) < n: Number of components to keep. :param normalize: (min, max) tuple of float: Optional range to normalize output values to. ''' n = array.shape[-1] all_axes_except_last = tuple(range(len(array.shape) - 1)) array_flat = array.reshape(-1, n) pca = sklearn.decomposition.PCA(n_components=k) pca.fit(array_flat) result_unnorm = pca.transform(array_flat).reshape(*array.shape[:-1], k) if normalize is not None: if result_unnorm.var() != 0.0: per_channel_min = result_unnorm.min( axis=all_axes_except_last, keepdims=True) per_channel_max = result_unnorm.max( axis=all_axes_except_last, keepdims=True) result = (result_unnorm - per_channel_min) / \ (per_channel_max - per_channel_min) result = result * (normalize[1] - normalize[0]) + normalize[0] else: result = np.ones_like(result_unnorm) * \ (normalize[0] + normalize[1]) / 2.0 else: result = result_unnorm result = result.astype(np.float32) return result @torch.no_grad() def cached_decode_latent_video_auto(my_dict, key, a4d_vae): ''' Ensures the video is only decoded once, stored in the dict, and retrieved automatically upon subsequent calls. :param key (str): rgb_* / depth_* / points_* / cams_* / ... :param latent_video: (B?, Cl, Tl, Hl, Wl) tensor of float. :return pixel_video: (B?, 1/3/6/9, Tp, Hp, Wp) tensor of float in [0, 1] or [-1, 1] (value range depends on modality). ''' if key + '_cdec' in my_dict: return my_dict[key + '_cdec'] latent_video = my_dict[key] # ^ (B?, Cl, Tl, Hl, Wl) tensor of float. has_batch_dim = (latent_video.ndim == 5) if not has_batch_dim: latent_video = latent_video[None] pixel_video = a4d_vae.decode_single_video_auto(key, latent_video) # ^ (B, 1/3/6/9, Tp, Hp, Wp) tensor of float in [0, 1] or [-1, 1]. if not has_batch_dim: pixel_video = pixel_video[0] my_dict[key + '_cdec'] = pixel_video # on GPU return my_dict[key + '_cdec'] @torch.no_grad() def cached_decode_latent_lowdim_auto(my_dict, key, a4d_vae): ''' Ensures the data is only decoded once, stored in the dict, and retrieved automatically upon subsequent calls. :param key (str): action / proprio / ... :param latent_data: (B?, T, C) tensor of float. :return raw_data: (B?, T, C) tensor of float. ''' if key + '_cdec' in my_dict: return my_dict[key + '_cdec'] latent_data = my_dict[key] # ^ (B?, T, C) tensor of float. has_batch_dim = (latent_data.ndim == 3) if not has_batch_dim: latent_data = latent_data[None] raw_data = a4d_vae.decode_single_lowdim_auto(key, latent_data) # ^ (B?, T, C) tensor of float. if not has_batch_dim: raw_data = raw_data[0] my_dict[key + '_cdec'] = raw_data # on GPU? return my_dict[key + '_cdec'] @torch.no_grad() def cached_decode_viz_latent_video_auto(my_dict, key, a4d_vae, config, first=True): ''' Ensures the video is only visualized once, stored in the dict, and retrieved automatically upon subsequent calls. :param latent_video (B?, Cl, Tl, Hl, Wl) tensor of float. :return viz_video (B?, 3, Tp, Hp, Wp) tensor of float in [-1, 1] (for consistency, value range is never [0, 1]). ''' if key + '_cviz' in my_dict: return my_dict[key + '_cviz'] pixel_video = cached_decode_latent_video_auto(my_dict, key, a4d_vae) # ^ (B?, 1/3/6/9, Tp, Hp, Wp) tensor of float in [0, 1] or [-1, 1]. has_batch_dim = (pixel_video.ndim == 5) if not has_batch_dim: pixel_video = pixel_video[None] elif first: pixel_video = pixel_video[0:1] viz_video = viz_video_auto(key, pixel_video, config) # ^ (B, 3, T, H, W) tensor of float in [-1, 1]. if has_batch_dim: viz_video = viz_video[0] my_dict[key + '_cviz'] = viz_video return viz_video @torch.no_grad() def viz_video_auto(key, pixel_video, config): ''' :param pixel_video (B?, 1/3/6/9, T, H, W) tensor of float in [0, 1] or [-1, 1] (value range depends on modality). :return viz_video (B?, 3, T, H, W) tensor of float in [-1, 1] (for consistency, value range is never [0, 1]). ''' has_batch_dim = (pixel_video.ndim == 5) if not has_batch_dim: pixel_video = pixel_video[None] (B, C, T, H, W) = pixel_video.shape assert B == 1, f'Expected B == 1, got {B}' if key.startswith('rgb'): viz_video = pixel_video * 2.0 - 1.0 elif key.startswith('depth'): if config.depth_viz_type == 'vidar': from externals.vidar.vidar.utils.viz import viz_depth video = [torch.tensor( viz_depth(video[0, :, i].float()), dtype=video.dtype, device=video.device).permute(2, 0, 1) for i in range(video.shape[2])] video = torch.stack(video, dim=1)[None] viz_video = video * 2.0 - 1.0 elif hasattr(plt.cm, config.depth_viz_type): video_np = pixel_video.float().cpu().numpy() # (B, 1, T, H, W) (B, _, T, H, W) = video_np.shape cmap = getattr(plt.cm, config.depth_viz_type) batch_frames = [] for t in range(T): frame = video_np[0, 0, t] # (H, W) # Apply colormap (returns RGBA, we take RGB) colored_frame = cmap(frame)[:, :, 0:3] # (H, W, 3) batch_frames.append(colored_frame) viz_video_np = np.stack(batch_frames, axis=0)[None] # (B, T, H, W, 3) viz_video_np = rearrange(viz_video_np, 'b t h w c -> b c t h w') # (B, 3, T, H, W) viz_video = torch.tensor(viz_video_np, dtype=pixel_video.dtype, device=pixel_video.device) viz_video = viz_video * 2.0 - 1.0 else: raise ValueError(f'Unknown or invalid depth_viz_type: {config.depth_viz_type}') elif key.startswith('points'): viz_video = pixel_video elif key.startswith('cams'): (video_rays, video_cross) = (pixel_video[:, 0:3], pixel_video[:, 3:6]) # NOTE(bvh): We deliberately never visualize the origin component because it would get too complex. if config.cams_viz_type == 'rays': viz_video = video_rays elif config.cams_viz_type == 'comb': # NOTE(bvh): I want to visualize rays and cross together here: # From left to right, we see columns: cross, rays, cross, rays, ... pixel_video_comb = video_rays.clone() cw = CAMS_COMB_WIDTH for i in range(cw // 2): pixel_video_comb[..., i::cw] = video_cross[..., i::cw] viz_video = pixel_video_comb else: raise ValueError(f'Unknown or invalid cams_viz_type: {config.cams_viz_type}') else: raise ValueError(f'Unknown or invalid sample key: {key}') if not has_batch_dim: viz_video = viz_video[0] return viz_video _REF_WIDTH = 480 # Reference width for text bar sizing (bar_height=24, font_scale=0.6 at this width). def text_bar_height(width): """Return text bar height scaled linearly with width (24px at 480px width).""" return max(int(round(width * 24 / _REF_WIDTH)), 10) def scaled_font(base_scale, width): """Return font_scale scaled linearly with width (base_scale at 480px width).""" return base_scale * width / _REF_WIDTH def paint_label(video, label, font_scale=0.6, font_color=(255, 255, 255), position='top_left', ref_width=None): ''' :param video: Input (T, H, W, 3) array of uint8. :param label (str): Text to apply onto the video. :param position (str): 'top_left', 'bottom_left', or 'bottom_right' :param ref_width: Width to use for font/bar scaling (defaults to video width). Pass a single-column width when painting labels on a full-gallery-width bar. :return video: Modified (T, H, W, 3) array of uint8. ''' # font = cv2.FONT_HERSHEY_SIMPLEX font = cv2.FONT_HERSHEY_DUPLEX # font = cv2.FONT_HERSHEY_TRIPLEX thickness = 1 H, W = video.shape[1], video.shape[2] # Scale font and bar height relative to ref_width (or actual width if not specified) scale_w = ref_width if ref_width is not None else W font_scale = scaled_font(font_scale, scale_w) * 0.9 bar_h = text_bar_height(scale_w) for i in range(video.shape[0]): frame = video[i] # Calculate text size text_size = cv2.getTextSize(label, font, font_scale, thickness)[0] text_w, text_h = text_size # Determine position if position == 'top_left': text_x = 0 text_y = bar_h rect_y1 = text_y - bar_h rect_y2 = text_y elif position == 'bottom_left': text_x = 0 text_y = H - 2 rect_y1 = H - bar_h rect_y2 = H elif position == 'top_right': text_x = W - text_w - 4 text_y = bar_h rect_y1 = text_y - bar_h rect_y2 = text_y elif position == 'bottom_right': text_x = W - text_w - 4 text_y = H - 2 rect_y1 = H - bar_h rect_y2 = H else: text_x = 0 text_y = bar_h rect_y1 = 0 rect_y2 = bar_h # Create a black rectangle with 50% transparency overlay = frame.copy() cv2.rectangle(overlay, (text_x, rect_y1), (text_x + text_w + 4, rect_y2), (0, 0, 0), -1) alpha = 0.5 # Transparency factor frame = cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0) # Put the text on the frame text_margin = max(int(round(6 * scale_w / _REF_WIDTH)), 2) frame = cv2.putText(frame, label, (text_x + 2, text_y - text_margin), font, font_scale, font_color, thickness, cv2.LINE_AA) video[i] = frame return video def annotate_mask_means(vids, mask_means): ''' :param vids: dict of dicts of (T, H, W, 3) arrays of uint8 with real keys. :param mask_means: dict of floats. :return vids: dict of dicts of (T, H, W, 3) arrays of uint8 with updated keys. ''' column_nature = { 'cond_x0': 'input', 'noisy_yt': 'output', 'assemb_xt': 'input', 'pred_y0': 'output', 'gt_y0': 'supervise', } new_vids = dict() for (column, nature) in column_nature.items(): new_vids[column] = dict() old_keys = list(vids[column].keys()) for k in old_keys: if k in mask_means[nature]: cur_mean = mask_means[nature][k] annotated_key = f'{k} ({cur_mean:.2f})' else: annotated_key = f'{k} (?)' new_vids[column][annotated_key] = vids[column].pop(k) return new_vids def make_gallery(vids, columns, headers, footers): ''' :param vids: dict of dicts of (T, H, W, 3) arrays of uint8. Per-tile decorations (mask indicator border, action / traj overlays, entry-key title, cam name / per-frame cam label, 'cams' tile height-subsample) are assumed to be baked in by the caller. Within each column, order of keys is preserved as to follow construction. :param columns: list of column keys to select from vids. :param headers: list of headers to add to the gallery. :param footers: list of footers to add to the gallery. :return gallery_vid: (T, H, W, 3) array of uint8. ''' col_vids = [] title_explain = { 'in': 'input', 'pred': 'prediction', 'gt': 'ground truth', } for col_key in columns: col_tiles = [] for (k, v) in vids[col_key].items(): assert v.dtype == np.uint8, f'Expected v.dtype == np.uint8, got {v.dtype}' col_tiles.append(v) # Pad tiles to the widest width within this column (centered) and to # the longest T (zero-pad at the end) so vertical concat is well-defined # even when entries have different frame counts. col_tiles = pad_videos_to_same_width(col_tiles, align='center') col_tiles = pad_videos_to_same_length(col_tiles) # Concatenate tiles vertically (variable heights per tile are fine). col_vid = np.concatenate(col_tiles, axis=1) # Write column title. label = title_explain[col_key] bh = text_bar_height(col_vid.shape[2]) col_title = paint_label(np.zeros_like(col_vid[:, 0:bh]), label) # Concatenate title on top. col_vid = np.concatenate([col_title, col_vid], axis=1) col_vids.append(col_vid) # Pad columns to the same height (top-aligned, pad at bottom) and to the # same T (zero-pad at the end) so horizontal concat is well-defined. col_vids = pad_videos_to_same_height(col_vids, align='top') col_vids = pad_videos_to_same_length(col_vids) # Concatenate columns horizontally. col_w = col_vids[0].shape[2] # single-column width for consistent text scaling gallery_vid = np.concatenate(col_vids, axis=2) # Concatenate headers on top. gbh = text_bar_height(col_w) for header in headers[::-1]: header_vid = paint_label(np.zeros_like(gallery_vid[:, 0:gbh]), header, font_color=(255, 255, 64), ref_width=col_w) gallery_vid = np.concatenate([header_vid, gallery_vid], axis=1) # Concatenate footers on bottom. for footer in footers: footer_vid = paint_label(np.zeros_like(gallery_vid[:, 0:gbh]), footer, font_color=(160, 160, 160), ref_width=col_w) gallery_vid = np.concatenate([gallery_vid, footer_vid], axis=1) assert gallery_vid.dtype == np.uint8, f'Expected gallery_vid.dtype == np.uint8, got {gallery_vid.dtype}' return gallery_vid def colormap_jet(val): ### TODO (Vitor): Replace with pyplot function ### Action colormap clr = [1.0, 1.0, 1.0] if val <= 0.33: clr[1] = val / 0.33 clr[0] = 0.0 elif (val > 0.33) & (val <= 0.67): clr[0] = (val - 0.33) / 0.33 clr[2] = 1.0 - clr[0] elif val > 0.67: clr[1] = 1.0 - (val - 0.67) / 0.33 clr[2] = 0.0 return [c * 2.0 - 1.0 for c in clr] @torch.no_grad() def unpack_nested(my_dict, key=None, batch_idx=0, dtype=str, optional=True): ''' Some simple values may get nested inside lists etc due to collating, so this method extracts them. NOTE: If there is a batch dimension, this only gets the FIRST sample. ''' if key is not None: value = my_dict.get(key, None) else: value = my_dict while isinstance(value, list) or (torch.is_tensor(value) and value.numel() > 1): value = value[batch_idx] if optional and value is None: return None if torch.is_tensor(value): value = value.item() return dtype(value) @torch.no_grad() def get_per_pixel_frame_mask_means(mask, a4d_vae): ''' :param mask: (1, Tl, Hl, Wl) tensor of float. :return frame_mm: (Tp) tensor of float containing averages for each raw frame. ''' T_ratio = a4d_vae.vae_ratio[0] Tl = mask.shape[1] Tp = (Tl - 1) * T_ratio + 1 frame_mm = torch.zeros(Tp, device=mask.device, dtype=torch.float32) for tl in range(Tl): tp1 = max((tl - 1) * T_ratio + 1, 0) tp2 = min(tl * T_ratio + 1, Tp) frame_mm[tp1:tp2] = mask[0, tl].mean() frame_mm = frame_mm.cpu() return frame_mm @torch.no_grad() def paint_mask_indicator(pixel_video, frame_mm, key, ios2rgb, border_width): ''' :param pixel_video: (3, Tp, Hp, Wp) tensor of float. Mutated in place. :param frame_mm: (Tp) tensor of float. :param ios2rgb: 3 x 3 matrix mapping (input, output, supervise) masks to (R, G, B) colors. ''' (C, Tp, Hp, Wp) = pixel_video.shape assert C == 3, f'Expected C == 3, got {C}' ios2rgb = torch.tensor(ios2rgb, dtype=torch.float32) # (3, 3) tensor of float. for t in range(Tp): # Dark gray indicates no conditioning or prediction is happening whatsoever. if frame_mm['input'][key][t] == 0.0 and frame_mm['output'][key][t] == 0.0: pixel_video[:, t] = 0.1 # Green border indicates conditioning is happening. # Blue border indicates prediction is happening. # Anything inbetween indicates spatial inpainting is happening. BW = border_width if BW > 0: border_color = ios2rgb @ torch.stack([ frame_mm['input'][key][t], frame_mm['output'][key][t], frame_mm['supervise'][key][t], ], dim=0) border_color = border_color.to(pixel_video.device, pixel_video.dtype) border_color = border_color[:, None, None] # (3, 1, 1) tensor of float. # NOTE(bvh): as of 01/15/2026, only top and bottom lines, # no vertical ones anymore, since ugly otherwise pixel_video[:, t, :BW, :] = border_color pixel_video[:, t, -BW:, :] = border_color def tiered_detail(iteration, max_detail, interval): ''' To frequently save fewer visuals, but occasionally save many visuals. ''' detail = 1 if max_detail >= 2 and iteration % (interval * 4) == 1: detail = 2 if max_detail >= 3 and iteration % (interval * 16) == 1: detail = 3 return detail def _natural_key(s): # Natural sort key so rgb10 sorts after rgb9, not after rgb1. return [int(t) if t.isdigit() else t for t in re.split(r'(\d+)', s)] def sorted_camslast(my_list): list1 = sorted([x for x in my_list if not x.startswith('cam')], key=_natural_key) list2 = sorted([x for x in my_list if x.startswith('cam')], key=_natural_key) sorted_list = list1 + list2 return sorted_list def _parse_align(align): '''Parse flexible alignment string into (before_frac, after_frac). 0.0 = top/left, 0.5 = center, 1.0 = bottom/right.''' a = align.lower().strip() if a in ('top', 'left', 'start'): return 0.0 elif a in ('center', 'mid', 'middle'): return 0.5 elif a in ('bottom', 'right', 'end'): return 1.0 return 0.5 # default center def pad_video_to_height(vid, target_h, align='center'): '''Pad a (T, H, W, C) video with black to target_h.''' h = vid.shape[1] if h >= target_h: return vid pad_total = target_h - h frac = _parse_align(align) pad_before = int(pad_total * frac) pad_after = pad_total - pad_before before = np.zeros((vid.shape[0], pad_before, *vid.shape[2:]), dtype=vid.dtype) after = np.zeros((vid.shape[0], pad_after, *vid.shape[2:]), dtype=vid.dtype) return np.concatenate([before, vid, after], axis=1) def pad_video_to_width(vid, target_w, align='center'): '''Pad a (T, H, W, C) video with black to target_w.''' w = vid.shape[2] if w >= target_w: return vid pad_total = target_w - w frac = _parse_align(align) pad_before = int(pad_total * frac) pad_after = pad_total - pad_before before = np.zeros((vid.shape[0], vid.shape[1], pad_before, vid.shape[3]), dtype=vid.dtype) after = np.zeros((vid.shape[0], vid.shape[1], pad_after, vid.shape[3]), dtype=vid.dtype) return np.concatenate([before, vid, after], axis=2) def pad_videos_to_same_height(vids_list, align='center'): '''Pad all (T, H, W, C) videos to the max height in the list.''' if len(vids_list) == 0: return vids_list max_h = max(v.shape[1] for v in vids_list) return [pad_video_to_height(v, max_h, align=align) for v in vids_list] def pad_videos_to_same_width(vids_list, align='center'): '''Pad all (T, H, W, C) videos to the max width in the list.''' if len(vids_list) == 0: return vids_list max_w = max(v.shape[2] for v in vids_list) return [pad_video_to_width(v, max_w, align=align) for v in vids_list] def pad_video_to_length(vid, target_t): '''Zero-pad a (T, H, W, C) video at the end (axis 0) to target_t.''' t = vid.shape[0] if t >= target_t: return vid pad = np.zeros((target_t - t, *vid.shape[1:]), dtype=vid.dtype) return np.concatenate([vid, pad], axis=0) def pad_videos_to_same_length(vids_list): '''Zero-pad all (T, H, W, C) videos at the end (axis 0) to the max T in the list.''' if len(vids_list) == 0: return vids_list max_t = max(v.shape[0] for v in vids_list) return [pad_video_to_length(v, max_t) for v in vids_list] def simple_horizontal_concat(vids, columns): ''' :param vids: dict of dicts of (T, H, W, 3) arrays of uint8. Within each column, we only take the first entry. :param columns: list of column keys to select from vids. :return concat_vid: (T, H_max, W_total, 3) array of uint8. ''' title_explain = { 'in': 'input', 'pred': 'prediction', 'gt': 'ground truth', } # Get first video from each column col_vids = [] for col_key in columns: first_video = next(iter(vids[col_key].values())) col_vids.append(first_video) # Pad columns to the same height (centered) and same T (zero-pad at end) col_vids = pad_videos_to_same_height(col_vids) col_vids = pad_videos_to_same_length(col_vids) # NOTE: shared bar height across columns to preserve the height alignment. bh = text_bar_height(max(v.shape[2] for v in col_vids)) for (col_idx, col_key) in enumerate(columns): col_vid = col_vids[col_idx] label = title_explain[col_key] col_title = paint_label(np.zeros_like(col_vid[:, 0:bh]), label) col_vids[col_idx] = np.concatenate([col_title, col_vid], axis=1) # Concatenate columns horizontally concat_vid = np.concatenate(col_vids, axis=2) return concat_vid def viz_custom(vids, entries, gt=False, headers=None): ''' Modular visualization: select specific entries (= columns), optionally show GT (= extra row), add headers (= on top). :param vids: dict of dicts of (T, H, W, 3) arrays of uint8. Expected keys: 'pred', 'gt', optionally 'in'. :param entries: list of entry keys to include (e.g., ['rgb1', 'rgb0', 'rgb2']). :param gt: whether to include GT row below prediction row. :param headers: list of header strings to add on top (or None). :return concat_vid: (T, H, W, 3) array of uint8. ''' if headers is None: headers = [] pred_vids = vids.get('pred', {}) gt_vids = vids.get('gt', {}) # Collect prediction videos in specified order pred_row = [pred_vids.get(k) for k in entries] pred_row = [v for v in pred_row if v is not None] if len(pred_row) == 0: return None # Concatenate prediction row horizontally (pad to equal height + T first) pred_row = pad_videos_to_same_height(pred_row) pred_row = pad_videos_to_same_length(pred_row) pred_concat = np.concatenate(pred_row, axis=2) # Optionally add GT row if gt: gt_row = [gt_vids.get(k) for k in entries] gt_row = [v for v in gt_row if v is not None] if len(gt_row) > 0: gt_row = pad_videos_to_same_height(gt_row) gt_row = pad_videos_to_same_length(gt_row) gt_concat = np.concatenate(gt_row, axis=2) # Resize GT row to match pred row width if needed if gt_concat.shape[2] != pred_concat.shape[2]: gt_concat = np.stack([cv2.resize( gt_concat[t], dsize=(pred_concat.shape[2], gt_concat.shape[1]), interpolation=cv2.INTER_CUBIC, ) for t in range(gt_concat.shape[0])], axis=0) # Match heights and T (centered H, zero-pad T at the end) so the # vertical concat below is well-defined. (pred_concat, gt_concat) = pad_videos_to_same_height( [pred_concat, gt_concat]) (pred_concat, gt_concat) = pad_videos_to_same_length( [pred_concat, gt_concat]) final_vid = np.concatenate([pred_concat, gt_concat], axis=1) else: final_vid = pred_concat else: final_vid = pred_concat # Add headers on top (use single-entry width for consistent text scaling) single_w = pred_row[0].shape[2] if len(pred_row) > 0 else final_vid.shape[2] hbh = text_bar_height(single_w) for header in headers[::-1]: header_vid = paint_label( np.zeros((final_vid.shape[0], hbh, final_vid.shape[2], 3), dtype=np.uint8), header, font_color=(255, 255, 64), ref_width=single_w) final_vid = np.concatenate([header_vid, final_vid], axis=1) return final_vid # def viz_anydrive_gaia(vids, entries): # ''' # Visualization for GAIA driving world model. # :param vids: dict of dicts of (T, H, W, 3) arrays of uint8. # Expected keys: 'pred'. # :param entries: list of entry keys to include (e.g., ['rgb1', 'rgb0', 'rgb2']). # :return concat_vid: (T, H, W, 3) array of uint8. # ''' # pred_vids = vids.get('pred', {}) # def get_video_by_entry_key(vids_dict, entry_key): # for k, v in vids_dict.items(): # if k.startswith(entry_key + ' ') or k == entry_key: # return v # return None # # Collect prediction videos in specified order # pred_row = [get_video_by_entry_key(pred_vids, k) for k in entries] # pred_row = [v for v in pred_row if v is not None] # if len(pred_row) == 0: # return None # # Concatenate prediction row horizontally (pad to equal height first) # pred_row = pad_videos_to_same_height(pred_row) # pred_concat = np.concatenate(pred_row, axis=2) # # TODO # final_vid = pred_concat # return final_vid