# Created by BVH, Jan 2026. # World model (action/trajectory) related visualization helpers. # Separated from visuals.py to keep file sizes manageable. from custom.dataloader import unified_anydrive from custom.eval.visuals import paint_label, pad_video_to_length import cv2 from imaginaire.utils import log from lovely_numpy import lo import lovely_tensors import matplotlib.pyplot as plt import numpy as np import torch def gripper_to_color(grip_value, arm_idx=0): ''' Convert gripper value to color based on arm identity. LEFT arm (arm_idx=1): Green (open) - Orange (closed) RIGHT arm (arm_idx=0): Cyan (open) - Orange (closed) :param grip_value: float in [0.0, 0.1], where 0.0=closed, 0.1=open. :param arm_idx: 0 for RIGHT arm (cyan), 1 for LEFT arm (green). :return: [R, G, B] list of floats in [0, 1]. ''' # Scale from [0, 0.1] to [0, 1] for interpolation (0=open, 1=closed) closedness = 1.0 - max(0.0, min(1.0, float(grip_value) * 10.0)) if arm_idx == 1: # LEFT arm: Green - Yellow r = 0.2 + 0.7 * closedness # 0.2 - 0.9 g = 1.0 - 0.3 * closedness # 1.0 - 0.7 b = 0.2 - 0.1 * closedness # 0.2 - 0.1 else: # RIGHT arm: Cyan - Purple r = 0.2 + 0.6 * closedness # 0.2 - 0.8 g = 0.8 - 0.4 * closedness # 0.8 - 0.4 b = 1.0 - 0.0 * closedness # 1.0 - 1.0 return [r, g, b] def fill_square_around_point(rgb, px, py, radius, color, bounds_safe=False): ''' Fill a square region around a point with a solid color. :param rgb: (3, H, W) tensor of float in [-1, 1]. :param px: int, x coordinate (column). :param py: int, y coordinate (row). :param radius: int, half-size of square (total size = 2*radius+1). :param color: list of 3 floats in [0, 1]. :param bounds_safe: If True, handle edge cases where square extends beyond image. :return rgb: (3, H, W) tensor of float in [-1, 1]. ''' radius = int(max(radius, 0)) H, W = rgb.shape[1], rgb.shape[2] if bounds_safe: # Compute clamped bounds y1 = max(0, py - radius) y2 = min(H, py + radius + 1) x1 = max(0, px - radius) x2 = min(W, px + radius + 1) if y2 <= y1 or x2 <= x1: return rgb # Nothing to draw # Create correctly sized square for the valid region h_size = y2 - y1 w_size = x2 - x1 square = torch.zeros((3, h_size, w_size), device=rgb.device, dtype=rgb.dtype) square[0] = color[0] * 2.0 - 1.0 square[1] = color[1] * 2.0 - 1.0 square[2] = color[2] * 2.0 - 1.0 rgb[:, y1:y2, x1:x2] = square else: # Original fast path (no bounds checking) square = torch.zeros((3, 2*radius+1, 2*radius+1), device=rgb.device, dtype=rgb.dtype) square[0] = color[0] * 2.0 - 1.0 square[1] = color[1] * 2.0 - 1.0 square[2] = color[2] * 2.0 - 1.0 rgb[:, py-radius:py+radius+1, px-radius:px+radius+1] = square return rgb def _clip_line_to_bounds(x0, y0, x1, y1, W, H): '''Cohen-Sutherland line clipping to [0, W-1] x [0, H-1]. Returns clipped (x0, y0, x1, y1) as ints, or None if fully outside.''' xmin, ymin, xmax, ymax = 0.0, 0.0, float(W - 1), float(H - 1) INSIDE, LEFT, RIGHT, BOTTOM, TOP = 0, 1, 2, 4, 8 x0, y0, x1, y1 = float(x0), float(y0), float(x1), float(y1) def outcode(x, y): code = INSIDE if x < xmin: code |= LEFT elif x > xmax: code |= RIGHT if y < ymin: code |= BOTTOM elif y > ymax: code |= TOP return code code0 = outcode(x0, y0) code1 = outcode(x1, y1) for _ in range(20): if not (code0 | code1): return (int(x0), int(y0), int(x1), int(y1)) if code0 & code1: return None code_out = code0 if code0 else code1 if code_out & TOP: x = x0 + (x1 - x0) * (ymax - y0) / (y1 - y0) y = ymax elif code_out & BOTTOM: x = x0 + (x1 - x0) * (ymin - y0) / (y1 - y0) y = ymin elif code_out & RIGHT: y = y0 + (y1 - y0) * (xmax - x0) / (x1 - x0) x = xmax elif code_out & LEFT: y = y0 + (y1 - y0) * (xmin - x0) / (x1 - x0) x = xmin if code_out == code0: x0, y0 = x, y code0 = outcode(x0, y0) else: x1, y1 = x, y code1 = outcode(x1, y1) return None def _draw_line(rgb, x0, y0, x1, y1, color, thickness=1): ''' Draw a line on a (3, H, W) tensor using Bresenham's algorithm. color: list of 3 floats in [0, 1]. ''' H, W = rgb.shape[1], rgb.shape[2] dx = abs(x1 - x0) dy = abs(y1 - y0) sx = 1 if x0 < x1 else -1 sy = 1 if y0 < y1 else -1 err = dx - dy color_vals = [c * 2.0 - 1.0 for c in color] half = thickness // 2 while True: for oy in range(-half, half + 1): for ox in range(-half, half + 1): py, px = y0 + oy, x0 + ox if 0 <= py < H and 0 <= px < W: for ch in range(3): rgb[ch, py, px] = color_vals[ch] if x0 == x1 and y0 == y1: break e2 = 2 * err if e2 > -dy: err -= dy x0 += sx if e2 < dx: err += dx y0 += sy return rgb @torch.no_grad() def project_action_arm(rgb, action, cam, arm_idx, condition=None, highlight_t=None, size=2): ''' Project a single robot arm's trajectory onto the image and draw it. Color indicates gripper state. Rotation axes are drawn as short lines from the EE position using the two rot6d vectors directly. :param rgb: (3, H, W) tensor of float in [-1, 1]. :param action: (1, T, 20) or (T, 20) tensor of action values. - arm_idx=0: RIGHT arm, indices 0:3 (xyz), 3:9 (rot6d), 18 (gripper) - arm_idx=1: LEFT arm, indices 9:12 (xyz), 12:18 (rot6d), 19 (gripper) :param cam: Camera object with project_points method and hw attribute. :param arm_idx: 0 for RIGHT arm, 1 for LEFT arm. :param condition: (1, T) or (T,) tensor indicating conditioned (1.0) vs predicted (0.0) frames. :param highlight_t: Frame index to highlight with gray border. :param size: Radius of drawn trajectory points. :return rgb: (3, H, W) tensor with arm trajectory drawn. ''' if action is None: return rgb action = action.clone().detach().float() if action.dim() == 3: action = action[0] # (T, 20) # Skip drawing if action contains NaN/Inf # TODO(bvh): not sure yet when or why this happens if not torch.isfinite(action).all(): log.error(f'NaN/Inf in action: {action}', rank0_only=False) action = action.clone() action[~torch.isfinite(action)] = 0.0 T_act = action.shape[0] # Extract per-arm fields: xyz (3), rot6d (6 = two 3D vectors), gripper (1) if arm_idx == 0: xyz = action[:, 0:3] # (T, 3) r6d = action[:, 3:9] # (T, 6) grip_values = action[:, 18] else: xyz = action[:, 9:12] r6d = action[:, 12:18] grip_values = action[:, 19] # NOTE(bvh): assume rot6d = [row0(3), row1(3)] already in world space. axis_len = 0.08 # 8cm ax0_ends = xyz + axis_len * r6d[:, 0:3] # (T, 3) ax1_ends = xyz + axis_len * r6d[:, 3:6] # (T, 3) # Stack all 3D points for batch projection: [origin, ax0_end, ax1_end] per timestep # cam.project_points expects (1, 3, N) all_pts = torch.stack([xyz, ax0_ends, ax1_ends], dim=1) # (T, 3, 3) all_pts = all_pts.reshape(-1, 3).unsqueeze(0).permute(0, 2, 1) # (1, 3, T*3) cam_device = cam.K.device if hasattr(cam, 'K') else rgb.device all_pts = all_pts.to(cam_device) # Project to image coordinates (normalized [-1, 1]) all_proj = cam.project_points(all_pts, from_world=True, flag_invalid=False) # (1, T*3, 2) all_proj = all_proj[0] # (T*3, 2) # Detect points behind camera using z in camera space pts_homo = torch.cat([all_pts, torch.ones(1, 1, all_pts.shape[2], device=all_pts.device)], dim=1) z_cam = torch.bmm(cam.Pwc(from_world=True).to(all_pts.device), pts_homo)[0, 2] # (T*3,) behind_cam = (z_cam <= 0).reshape(T_act, 3) # (T, 3): [origin, ax0_end, ax1_end] # Convert from normalized [-1, 1] to pixel coordinates H, W = int(cam.hw[0]), int(cam.hw[1]) all_proj_pix = all_proj.clone() all_proj_pix[:, 0] = (all_proj[:, 0] + 1) / 2 * W all_proj_pix[:, 1] = (all_proj[:, 1] + 1) / 2 * H # Reshape: (T, 3, 2) -> [origin, ax0_end, ax1_end] per timestep pts_pix = all_proj_pix.reshape(T_act, 3, 2) # Clamp origin to image bounds (axis endpoints left raw for smart clipping) pts_pix[:, 0, 0] = torch.clamp(pts_pix[:, 0, 0], 2, W - 3) pts_pix[:, 0, 1] = torch.clamp(pts_pix[:, 0, 1], 2, H - 3) max_axis_px = 0.3 * min(H, W) if condition is not None: condition = condition.clone().detach() if condition.dim() == 2: condition = condition[0] ax0_color = [1.0, 0.2, 0.2] # red ax1_color = [0.2, 0.2, 1.0] # blue def _clamp_axis_len(ox, oy, ax, ay): '''Shorten axis endpoint so projected length <= max_axis_px.''' dx, dy = ax - ox, ay - oy length = (dx * dx + dy * dy) ** 0.5 if length > max_axis_px: scale = max_axis_px / length ax = ox + dx * scale ay = oy + dy * scale return int(ax), int(ay) def _draw_timestep(j, extra_size=0, override_color=None, draw_axes=True): if behind_cam[j, 0]: return # Origin behind camera, skip entirely if not (torch.isfinite(pts_pix[j]).all()): return # NaN/Inf in projected coords, skip ox, oy = int(pts_pix[j, 0, 0]), int(pts_pix[j, 0, 1]) nonlocal rgb if draw_axes: # Axis 0 (red): skip if endpoint behind camera, else clip to image if not behind_cam[j, 1]: a0x, a0y = _clamp_axis_len(ox, oy, float(pts_pix[j, 1, 0]), float(pts_pix[j, 1, 1])) clipped = _clip_line_to_bounds(ox, oy, a0x, a0y, W, H) if clipped: rgb = _draw_line(rgb, *clipped, ax0_color) # Axis 1 (blue): skip if endpoint behind camera, else clip to image if not behind_cam[j, 2]: a1x, a1y = _clamp_axis_len(ox, oy, float(pts_pix[j, 2, 0]), float(pts_pix[j, 2, 1])) clipped = _clip_line_to_bounds(ox, oy, a1x, a1y, W, H) if clipped: rgb = _draw_line(rgb, *clipped, ax1_color) color = override_color if override_color else gripper_to_color(grip_values[j].item(), arm_idx) rgb = fill_square_around_point(rgb, ox, oy, size + extra_size, color, bounds_safe=True) # Draw non-highlighted timesteps first (axes only every 2nd frame) for j in range(T_act): if j == highlight_t: continue _draw_timestep(j, draw_axes=(j % 2 == 0)) # Draw highlighted frame last (on top) if highlight_t is not None and 0 <= highlight_t < T_act: _draw_timestep(highlight_t, extra_size=2, override_color=[0.9, 0.9, 0.9]) _draw_timestep(highlight_t, extra_size=0, override_color=[0.1, 0.1, 0.1]) return rgb def get_deviation_color(deviation, max_dev=0.4): ''' Interpolate color from green to orange based on deviation magnitude. :param deviation: Scalar deviation value (e.g., L2 norm of position difference). :param max_dev: Maximum deviation for full orange (default 0.4 meters). :return: (R, G, B) tuple of uint8 values. ''' alpha = min(deviation / max_dev, 1.0) # Green RGB: (100, 255, 100) -> Orange RGB: (255, 165, 0) r = int(100 + alpha * 155) # 100 -> 255 g = int(255 - alpha * 90) # 255 -> 165 b = int(100 - alpha * 100) # 100 -> 0 return (r, g, b) @torch.no_grad() def draw_action_coord_text(rgb, action, orig_action, t): ''' Draw 3D coordinates of both robot arms as text overlay on the image. Color matches trajectory squares based on gripper state: - RIGHT arm: Cyan (open) - Orange (closed) - LEFT arm: Green (open) - Orange (closed) :param rgb: (3, H, W) tensor of float in [-1, 1]. :param action: (1, T, 20) or (T, 20) tensor of action values to display. :param orig_action: Unused (kept for API compatibility). :param t: Current frame index. :return rgb: (3, H, W) tensor with coordinate text drawn. ''' # Normalize action shape to (T, 20) act = action[0] if action.dim() == 3 else action if t >= act.shape[0]: return rgb # Format coordinate text with gripper values # Indices 0,1,2,18 = RIGHT arm; indices 9,10,11,19 = LEFT arm xR, yR, zR = act[t, 0].item(), act[t, 1].item(), act[t, 2].item() xL, yL, zL = act[t, 9].item(), act[t, 10].item(), act[t, 11].item() gripR, gripL = act[t, 18].item(), act[t, 19].item() text1 = f'REE: ({xR:.2f}, {yR:.2f}, {zR:.2f}) g={gripR:.2f}' text2 = f'LEE: ({xL:.2f}, {yL:.2f}, {zL:.2f}) g={gripL:.2f}' # Get colors based on gripper state (same as trajectory squares) # RIGHT arm (0): Cyan-Orange, LEFT arm (1): Green-Orange color_R = gripper_to_color(gripR, arm_idx=0) # RIGHT arm color_L = gripper_to_color(gripL, arm_idx=1) # LEFT arm # Convert [0,1] to [0,255] for paint_label color1 = tuple(int(c * 255) for c in color_R) color2 = tuple(int(c * 255) for c in color_L) # Convert tensor [-1, 1] -> numpy [0, 255] uint8 device, dtype = rgb.device, rgb.dtype rgb_np = ((rgb.permute(1, 2, 0).cpu().float().numpy() + 1.0) / 2.0 * 255).astype(np.uint8) rgb_np = rgb_np[None] # (1, H, W, 3) for paint_label # Draw text labels with gripper-based colors (LEE on left, REE on right) paint_label(rgb_np, text2, font_scale=0.4, font_color=color2, position='bottom_left') paint_label(rgb_np, text1, font_scale=0.4, font_color=color1, position='bottom_right') # Convert back: numpy [0, 255] -> tensor [-1, 1] rgb_out = torch.from_numpy(rgb_np[0].astype(np.float32) / 255.0 * 2.0 - 1.0).permute(2, 0, 1) rgb_out = rgb_out.to(device=device, dtype=dtype) return rgb_out @torch.no_grad() def draw_actions(rgb, cam, action, proprio=None, condition=None, t=None, size=2, orig_action=None, show_coords=False): ''' Draw robot arm trajectories onto an RGB image with optional coordinate text. :param rgb: (3, H, W) tensor of float in [-1, 1]. :param cam: Camera object for 3D-to-2D projection. :param action: (1, T, 20) or (T, 20) tensor of action values (both arms). :param proprio: (1, T, 20) optional proprioception tensor (drawn in white). :param condition: (1, T) tensor indicating conditioned (1.0) vs predicted (0.0) frames. :param t: Current frame index to highlight (drawn in black). :param size: Radius of drawn trajectory points. :param orig_action: (1, T, 20) original/clean action for deviation-based text color. :param show_coords: If True, draw 3D coordinates as text in bottom corners. :return rgb: (3, H, W) tensor with trajectories and optional text drawn. ''' rgb = rgb.clone().detach() # Draw action trajectories for both arms (arm 0 = left, arm 1 = right) rgb = project_action_arm(rgb, action, cam, arm_idx=0, condition=condition, highlight_t=t, size=size) rgb = project_action_arm(rgb, action, cam, arm_idx=1, condition=condition, highlight_t=t, size=size) # Draw proprio trajectories in white if provided if proprio is not None: rgb = project_action_arm(rgb, proprio, cam, arm_idx=0, condition=None, highlight_t=None, size=size) rgb = project_action_arm(rgb, proprio, cam, arm_idx=1, condition=None, highlight_t=None, size=size) # Draw 3D coordinate text overlay if requested if show_coords and action is not None and t is not None: rgb = draw_action_coord_text(rgb, action, orig_action, t) return rgb @torch.no_grad() def get_traj_bounds(traj1, traj2, traj3): ''' :param traj*: (T, 2) tensor of float with (x, -z) values for each timestep. :return traj_bounds: (4) tuple of float: (traj_max, traj_min, traj_range, traj_norm). ''' fallback = traj1 if traj1 is not None else traj2 if traj2 is not None else traj3 traj1 = traj1 if traj1 is not None else fallback traj2 = traj2 if traj2 is not None else fallback traj3 = traj3 if traj3 is not None else fallback traj1 = torch.nan_to_num(traj1.clone().detach(), 0.0) traj2 = torch.nan_to_num(traj2.clone().detach(), 0.0) traj3 = torch.nan_to_num(traj3.clone().detach(), 0.0) # Use forward (Z) coordinate only for normalization # since it is often much larger than lateral (X) coordinate. epsilon = torch.tensor(1e-3) chosen_max = max(traj1.max(), traj2.max(), traj3.max()).item() chosen_min = min(traj1.min(), traj2.min(), traj3.min()).item() chosen_range = max(chosen_max - chosen_min, epsilon) chosen_norm = max(traj1.abs().max(), traj2.abs().max(), traj3.abs().max(), epsilon).item() return (chosen_max, chosen_min, chosen_range, chosen_norm) @torch.no_grad() def draw_traj(rgb, traj, condition, is_clean, highlight, traj_bounds, size, show_coords=True): ''' :param rgb: (3, H, W) tensor of float. :param traj: (T, 2) tensor of float. :param condition: (T) tensor of float. :param highlight: int. :param traj_bounds: (4) tuple of float. :param size: int. :param show_coords: bool, whether to show EGO coordinates in lower left. :return rgb: (3, H, W) tensor of float. ''' viz_rgb = rgb.clone().detach() traj = traj.clone().detach() # condition = condition.clone().detach() T = traj.shape[0] (H, W) = rgb.shape[1:3] # traj_max = traj.max() # traj_min = traj.min() # traj_range = traj_max - traj_min # traj_norm = traj.abs().max() (traj_max, traj_min, traj_range, traj_norm) = traj_bounds for t in range(T): # tx = positive right, tz = positive forward (tx, tz) = traj[t] tx = float(torch.nan_to_num(tx, 0.0)) tz = float(torch.nan_to_num(tz, 0.0)) tx = (tx / traj_norm + 1.0) / 2.0 # in [0, 1] tz = (tz - traj_min) / traj_range # in [0, 1] px = int(tx * (W - 20)) + 10 py = H - (int(tz * (H - 20)) + 10) if t == highlight: viz_rgb = fill_square_around_point(viz_rgb, px, py, size + 1, [0.8, 0.8, 0.8]) rem_size = size - 1 else: rem_size = size if is_clean: color = [0.2, 0.2, 0.2] # the original GT traj elif condition[t] != 0.0: color = plt.cm.plasma(t / T) # our custom input traj else: color = np.array(plt.cm.plasma(t / T)) * 0.4 + 0.4 # our custom predicted traj viz_rgb = fill_square_around_point(viz_rgb, px, py, rem_size, color) # Draw EGO coordinates text in lower left corner if show_coords: ego_x, ego_z = traj[highlight].tolist() viz_rgb = draw_traj_coord_text(viz_rgb, ego_x, ego_z) return viz_rgb def draw_traj_coord_text(rgb, ego_x, ego_z): ''' Draw EGO coordinates text in lower left corner. :param rgb: (3, H, W) tensor of float in [-1, 1]. :param ego_x: float, x coordinate. :param ego_z: float, z coordinate. :return rgb: (3, H, W) tensor of float in [-1, 1]. ''' label = f'EGO: ({ego_x:.2f}, {ego_z:.2f})' font_color = (255, 255, 64) # yellow # Convert tensor [-1, 1] -> numpy [0, 255] uint8 device, dtype = rgb.device, rgb.dtype rgb_np = ((rgb.permute(1, 2, 0).cpu().float().numpy() + 1.0) / 2.0 * 255).astype(np.uint8) rgb_np = rgb_np[None] # (1, H, W, 3) for paint_label # Draw text label paint_label(rgb_np, label, font_scale=0.5, font_color=font_color, position='bottom_left') # Convert back: numpy [0, 255] -> tensor [-1, 1] rgb_out = torch.from_numpy(rgb_np[0].astype(np.float32) / 255.0 * 2.0 - 1.0).permute(2, 0, 1) rgb_out = rgb_out.to(device=device, dtype=dtype) return rgb_out _GAIA_CAM_ORDER = unified_anydrive.CAM_ORDER['gaia'] # ^ corresponds to rgb0, rgb1, ..., rgb4 (or up to rgb10). def viz_anydrive_gaia(pred_vids, sel_cam_names=None): ''' This is basically a more advanced, tailored version of viz_custom() that stacks 5-camera or 11-camera configurations in a predetermined 2D pattern. Detects layout variant (gaia11, gaia5-pvm, gaia5-wide) from cam names. :param pred_vids: dict of (T, H, W, 3) arrays of uint8. Expected keys: 'rgb0*', 'rgb1*', ..., 'rgb4*' or 'rgb10*'. :param sel_cam_names: list of cam names indexed by view (rgb0, rgb1, ...). If None, falls back to _GAIA_CAM_ORDER prefix. :return concat_vid: (T, H, W, 3) array of uint8. ''' # NOTE(bvh): Layout aspect ratios are approximate; frames are resized to fit each cell. num_views = len(pred_vids) if sel_cam_names is None: sel_cam_names = _GAIA_CAM_ORDER[:num_views] # Map camera name to view index (for looking up rgb from a cam name). cam_to_view = {} for v_idx, cam in enumerate(sel_cam_names): cam_to_view.setdefault(cam, v_idx) pvm_cams = {'front_pvm', 'left_side_pvm', 'right_side_pvm', 'rear_pvm'} wide_cams = {'left_side_forward', 'left_side_rearward', 'right_side_forward', 'right_side_rearward'} tele_cams = {'front_tele', 'rear_tele'} cam_set = set(sel_cam_names) has_pvm = pvm_cams.issubset(cam_set) has_wide = wide_cams.issubset(cam_set) has_tele = tele_cams.issubset(cam_set) # Layout cells: (cam_name, x, y, w, h) in arbitrary layout units. if has_pvm and has_wide and has_tele: # gaia11: full layout, see GAIA11_Viz1.png. cells = [ ('left_side_forward', 0, 0, 90, 50), ('left_side_rearward', 0, 50, 90, 50), ('front_tele', 0, 100, 90, 50), ('front_wide', 90, 0, 200, 110), ('front_pvm', 90, 110, 50, 40), ('left_side_pvm', 140, 110, 50, 40), ('right_side_pvm', 190, 110, 50, 40), ('rear_pvm', 240, 110, 50, 40), ('right_side_forward', 290, 0, 90, 50), ('right_side_rearward',290, 50, 90, 50), ('rear_tele', 290, 100, 90, 50), ] (layout_w, layout_h) = (380, 150) elif has_pvm: # gaia5-pvm: front_wide on top, 4 PVM cells underneath. cells = [ ('front_wide', 0, 0, 200, 110), ('front_pvm', 0, 110, 50, 40), ('left_side_pvm', 50, 110, 50, 40), ('right_side_pvm', 100, 110, 50, 40), ('rear_pvm', 150, 110, 50, 40), ] (layout_w, layout_h) = (200, 150) elif has_wide: # gaia5-wide: front_wide center, forward/rearward stacked on each side. cells = [ ('left_side_forward', 0, 0, 90, 50), ('left_side_rearward', 0, 50, 90, 50), ('front_wide', 90, 0, 180, 100), ('right_side_forward', 270, 0, 90, 50), ('right_side_rearward',270, 50, 90, 50), ] (layout_w, layout_h) = (360, 100) else: log.warning(f'viz_anydrive_gaia: could not detect layout variant for ' f'cams={sel_cam_names}; skipping.', rank0_only=True) return None # Scale so the front_wide cell becomes 2x the rgb0 (= front_wide) resolution. # Example canvas shapes (T, H, W, 3) for rgb0_w=384 (usual setting): # gaia11 (T, 576, 1459, 3) # gaia5-pvm (T, 576, 768, 3) # gaia5-wide (T, 427, 1536, 3) fw_v_idx = cam_to_view.get('front_wide') fw_vid = pred_vids.get(f'rgb{fw_v_idx}') if fw_v_idx is not None else None if fw_vid is None: log.warning(f'viz_anydrive_gaia: front_wide not in pred_vids; skipping.', rank0_only=True) return None fw_cell_w = next(w for (n, _x, _y, w, _h) in cells if n == 'front_wide') # Canvas T = max across all participating cell videos so views with shorter # frame counts get zero-padded at the end rather than crashing the assign. cell_vids = [] for (cam_name, _x, _y, _w, _h) in cells: v_idx = cam_to_view.get(cam_name) if v_idx is None: continue vid = pred_vids.get(f'rgb{v_idx}') if vid is not None: cell_vids.append(vid) T = max(v.shape[0] for v in cell_vids) if cell_vids else fw_vid.shape[0] scale = (2.0 * fw_vid.shape[2]) / fw_cell_w canvas_h = int(round(layout_h * scale)) canvas_w = int(round(layout_w * scale)) canvas = np.zeros((T, canvas_h, canvas_w, 3), dtype=np.uint8) for (cam_name, x, y, w, h) in cells: v_idx = cam_to_view.get(cam_name) if v_idx is None: log.warning(f'viz_anydrive_gaia: cam {cam_name!r} not in sel_cam_names; ' f'skipping cell.', rank0_only=True) continue vid = pred_vids.get(f'rgb{v_idx}') if vid is None: continue vid = pad_video_to_length(vid, T) cell_x0 = int(round(x * scale)) cell_y0 = int(round(y * scale)) cell_x1 = int(round((x + w) * scale)) cell_y1 = int(round((y + h) * scale)) cell_w = cell_x1 - cell_x0 cell_h = cell_y1 - cell_y0 resized = np.stack([cv2.resize( vid[t], dsize=(cell_w, cell_h), interpolation=cv2.INTER_AREA, ) for t in range(T)], axis=0) canvas[:, cell_y0:cell_y1, cell_x0:cell_x1] = resized return canvas