''' BVH, Nov 2025. Gradio utilities for Any4D / AnyView. Includes camera visualization for input/output videos. POSE REPRESENTATION SYSTEM: --------------------------- There are three pose representations used in this codebase: 1. ABSOLUTE CAM2WORLD (canonical representation): - Used for: visualization, tarfiles (by default), canonical poses - Format: 4x4 matrix transforming from camera space to world space - All poses are independent and absolute 2. RELATIVE WORLD2CAM (vidar internal format): - Used for: vidar['pose'] dictionary - Format: 4x4 matrix transforming from world space to camera space - Two-level hierarchy: * (0, 0): ABSOLUTE world2cam for base camera/frame * (0, c) for c != 0: RELATIVE to (0, 0) absolute_world2cam(0,c) = (0,c)_relative @ (0,0) * (t, c) for t != 0: RELATIVE to (0, c) of same camera absolute_world2cam(t,c) = (t,c)_relative @ (0,c) - This relative format allows efficient batch processing 3. ABSOLUTE WORLD2CAM (tarfiles with Invert checkbox): - Used for: tarfiles when "Invert" checkbox is checked - Format: 4x4 matrix transforming from world space to camera space - All poses are independent and absolute - Requires inversion to convert to CAM2WORLD before use Conversion functions: - convert_absolute_to_vidar_format: ABSOLUTE CAM2WORLD -> RELATIVE WORLD2CAM - convert_vidar_to_absolute_format: RELATIVE WORLD2CAM -> ABSOLUTE CAM2WORLD - invert_tarfile_poses: ABSOLUTE WORLD2CAM -> ABSOLUTE CAM2WORLD ''' import copy import numpy as np import plotly.graph_objects as go import torch import yaml from rich import print def create_canonical_camera_poses(azimuth_deg=0, input_elevation_deg=0, output_elevation_deg=0, radius=5.0, trajectory_type='Jumpy', num_timesteps=41, which_setting='Exo2Exo', timestep_indices=None): ''' Create canonical camera poses based on azimuth and elevation angles. Returns ABSOLUTE CAM2WORLD poses. Elevation is specified as absolute angles from the horizon (0° = horizon, 90° = straight up). Input camera has specified azimuth (usually 0) and elevation. Output camera trajectory depends on trajectory_type: - 'Jumpy': Output camera is static at the destination position (azimuth + elevation) - 'Gradual': Output camera interpolates from input to output angles over time in spherical coordinates (linear) - 'GradualSqrt': Output camera interpolates using sqrt(t) for smoother acceleration at start :param azimuth_deg (float): Azimuth rotation angle in degrees (negative = left, positive = right) :param input_elevation_deg (float): Input camera elevation in degrees from horizon (0 = horizon, 90 = up, -90 = down) :param output_elevation_deg (float): Output camera elevation in degrees from horizon :param radius (float): Distance from center in meters :param trajectory_type (str): 'Jumpy', 'Gradual', or 'GradualSqrt' :param num_timesteps (int): Number of timesteps :param which_setting (str): 'Exo2Exo', 'Ego2Ego', or 'Ego2Exo' :return: dict with (t, c) -> (1, 4, 4) ABSOLUTE CAM2WORLD pose matrices ''' if abs(radius) < 1e-6: radius = 1e-6 # Handle Ego2Ego setting: flip radius and invert angles if which_setting == 'Ego2Ego': radius = -radius azimuth_deg = -azimuth_deg input_elevation_deg = -input_elevation_deg output_elevation_deg = -output_elevation_deg # Convert to radians azimuth_rad = np.deg2rad(azimuth_deg) input_elevation_rad = np.deg2rad(input_elevation_deg) output_elevation_rad = np.deg2rad(output_elevation_deg) # Input camera: at specified elevation (azimuth=0 for input) input_azimuth = 0.0 input_elevation = input_elevation_rad # Output camera: at specified azimuth and elevation target_azimuth = azimuth_rad target_elevation = output_elevation_rad def create_camera_pose(azimuth, elevation, radius): ''' Create 4x4 camera-to-world pose matrix (extrinsic). Camera convention: X right, Y down, Z forward World convention: Z axis up ''' # Camera position in world coordinates (spherical to cartesian) # World Z is up, so elevation affects Z coordinate x = radius * np.cos(elevation) * np.cos(azimuth) y = radius * np.cos(elevation) * np.sin(azimuth) z = radius * np.sin(elevation) position = np.array([x, y, z]) # Camera Z axis (forward): # - For Exo2Exo/Ego2Exo: points towards origin (exocentric view) # - For Ego2Ego: points away from origin (egocentric view) if which_setting == 'Ego2Ego': forward = position / np.linalg.norm(position) else: forward = -position / np.linalg.norm(position) # World up vector (Z axis up) world_up = np.array([0, 0, 1]) # Camera X axis (right): perpendicular to forward and world up # right = forward × world_up (cross product) right = np.cross(forward, world_up) right = right / (np.linalg.norm(right) + 1e-8) # Camera Y axis (down): perpendicular to forward and right # For right-handed system: right × down = forward # Therefore: down = forward × right down = np.cross(forward, right) down = down / (np.linalg.norm(down) + 1e-8) # Build rotation matrix with camera axes as columns # R = [X_cam, Y_cam, Z_cam] = [right, down, forward] R = np.column_stack([right, down, forward]) # Build 4x4 pose matrix (camera-to-world transform) pose = np.eye(4) pose[:3, :3] = R pose[:3, 3] = position return torch.from_numpy(pose).float() def create_ego_camera_pose(): ''' Create 4x4 camera-to-world pose matrix for ego camera at origin. Camera convention: X right, Y down, Z forward Looking horizontally forward (in world where Z is up, camera looks along -X) ''' # Camera at origin position = np.array([0.0, 0.0, 0.0]) # Camera axes in world coordinates # World Z is up, so camera should look horizontally (along -X) right = np.array([0.0, 1.0, 0.0]) # X axis (right in camera = +Y in world) down = np.array([0.0, 0.0, -1.0]) # Y axis (down in world = -Z) forward = np.array([-1.0, 0.0, 0.0]) # Z axis (forward in camera = -X in world) # Build rotation matrix with camera axes as columns R = np.column_stack([right, down, forward]) # Build 4x4 pose matrix pose = np.eye(4) pose[:3, :3] = R pose[:3, 3] = position return torch.from_numpy(pose).float() # Create input pose (static) if which_setting == 'Ego2Exo': # Input camera at origin input_pose = create_ego_camera_pose() else: # Standard Exo2Exo or Ego2Ego input_pose = create_camera_pose(input_azimuth, input_elevation, radius) # Build pose dictionary (t, c) -> pose # VIDAR CONVENTION: Camera 0 = output/target, Camera 1 = input/context pose_dict = {} # Use provided timestep_indices if available, otherwise use range(num_timesteps) if timestep_indices is not None: time_indices = sorted(timestep_indices) else: time_indices = list(range(num_timesteps)) if trajectory_type == 'Jumpy': # Output camera is static at destination position output_pose = create_camera_pose(target_azimuth, target_elevation, radius) for t in time_indices: pose_dict[(t, 0)] = output_pose.unsqueeze(0) # Camera 0 = output (static) pose_dict[(t, 1)] = input_pose.unsqueeze(0) # Camera 1 = input (static) elif trajectory_type in ['Gradual', 'GradualSqrt']: # Output camera interpolates from input to target angles over time # Gradual uses linear interpolation, GradualSqrt uses sqrt for smoother acceleration for idx, t in enumerate(time_indices): # Compute interpolation factor based on trajectory type t_normalized = idx / max(1, len(time_indices) - 1) if trajectory_type == 'Gradual': alpha = t_normalized # Linear interpolation else: # GradualSqrt alpha = np.sqrt(t_normalized) # Sqrt interpolation for smoother acceleration if which_setting == 'Ego2Exo': # For Ego2Exo, interpolate from origin (ego camera) to target position in spherical coordinates # Interpolate spherical coordinates (azimuth, elevation, radius) interp_azimuth = 0.0 + alpha * target_azimuth interp_elevation = 0.0 + alpha * target_elevation interp_radius = 0.0 + alpha * radius # Radius interpolates from 0 (ego at origin) to target radius # Create pose using interpolated spherical coordinates # This ensures the camera follows the spherical trajectory if abs(interp_radius) < 1e-6: # When radius is near zero, use ego camera pose output_pose = input_pose.clone() else: # Create camera pose at interpolated spherical position output_pose = create_camera_pose(interp_azimuth, interp_elevation, interp_radius) else: # Standard interpolation for Exo2Exo and Ego2Ego # Interpolated angles interp_azimuth = input_azimuth + alpha * (target_azimuth - input_azimuth) interp_elevation = input_elevation + alpha * (target_elevation - input_elevation) # Create interpolated output pose output_pose = create_camera_pose(interp_azimuth, interp_elevation, radius) pose_dict[(t, 0)] = output_pose.unsqueeze(0) # Camera 0 = output (interpolated) pose_dict[(t, 1)] = input_pose.unsqueeze(0) # Camera 1 = input (static) else: raise ValueError(f"Unknown trajectory_type: {trajectory_type}. Must be 'Jumpy', 'Gradual', or 'GradualSqrt'") return pose_dict def convert_absolute_to_vidar_format(absolute_poses): ''' Convert absolute cam2world poses to vidar storage format. Vidar format (RELATIVE WORLD2CAM) - Two-level hierarchy: - (0, 0): ABSOLUTE world2cam for base camera/frame - (0, c) for c != 0: RELATIVE to (0, 0), i.e., absolute_world2cam(0,c) = (0,c)_relative @ (0,0) - (t, c) for t != 0: RELATIVE to (0, c), i.e., absolute_world2cam(t,c) = (t,c)_relative @ (0,c) :param absolute_poses (dict): Dict of (t, c) -> (B, 4, 4) absolute cam2world poses :return (dict): Dict of (t, c) -> (B, 4, 4) poses in vidar RELATIVE WORLD2CAM format ''' vidar_poses = {} # First, convert all cam2world poses to world2cam by inverting world2cam_absolute = {} for key, pose_cam2world in absolute_poses.items(): world2cam_absolute[key] = torch.inverse(pose_cam2world) # Get the base pose (0, 0) - absolute world2cam key_0_0 = (0, 0) if key_0_0 not in world2cam_absolute: raise ValueError(f"Base pose (0, 0) not found in absolute_poses") pose_0_0_world2cam = world2cam_absolute[key_0_0] # (B, 4, 4) vidar_poses[key_0_0] = pose_0_0_world2cam # Compute inverse of base pose for camera-level relative computation pose_0_0_world2cam_inv = torch.inverse(pose_0_0_world2cam) # (B, 4, 4) # Find all unique camera indices all_cameras = set([key[1] for key in world2cam_absolute.keys()]) # Store absolute poses for each camera at t=0 (needed for temporal relatives) camera_0_absolute = {} # Maps camera_idx -> absolute world2cam at t=0 # Process all cameras at t=0 first (convert to relative to (0,0)) for cam_idx in all_cameras: key_0_c = (0, cam_idx) if key_0_c in world2cam_absolute: pose_0_c_world2cam = world2cam_absolute[key_0_c] camera_0_absolute[cam_idx] = pose_0_c_world2cam if cam_idx == 0: # (0, 0) already stored above continue else: # (0, c) is relative to (0, 0) pose_relative = torch.matmul(pose_0_c_world2cam, pose_0_0_world2cam_inv) vidar_poses[key_0_c] = pose_relative # Process all frames at t>0 (convert to relative to their camera's t=0) for key, pose_world2cam in world2cam_absolute.items(): t, c = key if t != 0: # (t, c) is relative to (0, c) if c not in camera_0_absolute: raise ValueError(f"Camera {c} at t=0 not found for key {key}") pose_0_c_world2cam = camera_0_absolute[c] pose_0_c_world2cam_inv = torch.inverse(pose_0_c_world2cam) # Compute relative: world2cam_(t,c) @ world2cam_(0,c)^-1 pose_relative = torch.matmul(pose_world2cam, pose_0_c_world2cam_inv) vidar_poses[key] = pose_relative return vidar_poses def convert_vidar_to_absolute_format(vidar_poses): ''' Convert vidar storage format to absolute cam2world poses. Vidar format (RELATIVE WORLD2CAM) - Two-level hierarchy: - (0, 0): ABSOLUTE world2cam for base camera/frame - (0, c) for c != 0: RELATIVE to (0, 0), i.e., absolute_world2cam(0,c) = (0,c)_relative @ (0,0) - (t, c) for t != 0: RELATIVE to (0, c), i.e., absolute_world2cam(t,c) = (t,c)_relative @ (0,c) :param vidar_poses (dict): Dict of (t, c) -> (B, 4, 4) poses in vidar RELATIVE WORLD2CAM format :return (dict): Dict of (t, c) -> (B, 4, 4) absolute cam2world poses ''' absolute_poses = {} # Get the base pose (0, 0) - absolute world2cam key_0_0 = (0, 0) if key_0_0 not in vidar_poses: raise ValueError(f"Base pose (0, 0) not found in vidar_poses") pose_0_0_world2cam = vidar_poses[key_0_0] # (B, 4, 4) # Convert base pose (0, 0) to cam2world by inverting pose_0_0_cam2world = torch.inverse(pose_0_0_world2cam) absolute_poses[key_0_0] = pose_0_0_cam2world # Find all unique camera indices all_cameras = set([key[1] for key in vidar_poses.keys()]) # Store absolute world2cam for each camera at t=0 (needed for temporal conversion) camera_0_world2cam = {} # Maps camera_idx -> absolute world2cam at t=0 camera_0_world2cam[0] = pose_0_0_world2cam # Process all cameras at t=0 first (convert from relative to absolute) for cam_idx in all_cameras: if cam_idx == 0: continue # Already processed key_0_c = (0, cam_idx) if key_0_c in vidar_poses: pose_0_c_relative = vidar_poses[key_0_c] # (0, c) is relative to (0, 0), so: absolute = relative @ (0,0) pose_0_c_world2cam = torch.matmul(pose_0_c_relative, pose_0_0_world2cam) camera_0_world2cam[cam_idx] = pose_0_c_world2cam # Convert to cam2world by inverting pose_0_c_cam2world = torch.inverse(pose_0_c_world2cam) absolute_poses[key_0_c] = pose_0_c_cam2world # Process all frames at t>0 (convert from relative to absolute) for key, pose_relative in vidar_poses.items(): t, c = key if t != 0: # (t, c) is relative to (0, c) if c not in camera_0_world2cam: raise ValueError(f"Camera {c} at t=0 not found for key {key}") pose_0_c_world2cam = camera_0_world2cam[c] # Compute absolute: world2cam_(t,c) = relative @ world2cam_(0,c) pose_world2cam = torch.matmul(pose_relative, pose_0_c_world2cam) # Convert to cam2world by inverting pose_cam2world = torch.inverse(pose_world2cam) absolute_poses[key] = pose_cam2world return absolute_poses def apply_canonical_trajectory(base_data_batch, azimuth_deg=0, input_elevation_deg=0, output_elevation_deg=0, radius=5.0, trajectory_type='Jumpy', which_setting='Exo2Exo'): ''' Apply canonical camera trajectory to base_data_batch. Flow: 1. Create canonical poses in ABSOLUTE CAM2WORLD format 2. Convert to RELATIVE WORLD2CAM (vidar format) 3. Store in data_batch['anydata']['pose'] :param base_data_batch (dict): Data batch to modify :param azimuth_deg (float): Azimuth rotation angle in degrees :param input_elevation_deg (float): Input camera elevation in degrees from horizon :param output_elevation_deg (float): Output camera elevation in degrees from horizon :param radius (float): Radius in meters :param trajectory_type (str): 'Jumpy', 'Gradual', or 'GradualSqrt' :param which_setting (str): 'Exo2Exo' or 'Ego2Ego' :return: Modified data batch with poses in RELATIVE WORLD2CAM format ''' # Make a copy to avoid modifying the original data_batch = copy.deepcopy(base_data_batch) # Get actual timestep indices from existing RGB data (most reliable source) actual_timesteps = None if 'anydata' in data_batch and 'rgb' in data_batch['anydata']: # Get unique timestep indices from RGB keys actual_timesteps = set([key[0] for key in data_batch['anydata']['rgb'].keys()]) num_actual_frames = len(actual_timesteps) max_timestep = max(actual_timesteps) if actual_timesteps else 0 print(f'[cyan]Detected {num_actual_frames} frames in batch with timestep indices: 0-{max_timestep}') elif 'anydata' in data_batch and 'timestep' in data_batch['anydata']: actual_timesteps = set([key[0] for key in data_batch['anydata']['timestep'].keys()]) num_actual_frames = len(actual_timesteps) print(f'[cyan]Detected {num_actual_frames} frames from timestep dict') else: num_actual_frames = 41 # Default print(f'[yellow]No RGB or timestep data found, using default {num_actual_frames} frames') # Create canonical poses (absolute cam2world) # Only create poses for timesteps that actually have RGB data canonical_poses = create_canonical_camera_poses( azimuth_deg=azimuth_deg, input_elevation_deg=input_elevation_deg, output_elevation_deg=output_elevation_deg, radius=radius, trajectory_type=trajectory_type, num_timesteps=num_actual_frames, which_setting=which_setting, timestep_indices=actual_timesteps ) # Convert to vidar format: frame 0 is absolute, frames > 0 are relative to frame 0 (per camera) vidar_poses = convert_absolute_to_vidar_format(canonical_poses) # Replace poses in vidar dict if 'anydata' in data_batch: data_batch['anydata']['pose'] = vidar_poses print(f'[green]Applied canonical trajectory with {num_actual_frames} frames: {trajectory_type}, setting={which_setting}, ' f'azimuth={azimuth_deg}°, input_elev={input_elevation_deg}°, output_elev={output_elevation_deg}°, radius={radius}m') return data_batch def invert_tarfile_poses(tarfile_data): ''' Invert camera poses in tarfile data (before converting to VidarDataset format). By default, tarfiles store ABSOLUTE CAM2WORLD poses. When the "Invert" checkbox is checked, the tarfile actually contains ABSOLUTE WORLD2CAM poses, so we need to convert them to CAM2WORLD by inverting. For a 4x4 camera pose matrix [R|t], the inverse is: [R^T | -R^T @ t] [ 0 | 1 ] :param tarfile_data (dict): Tarfile data dict with 'pose' key (ABSOLUTE WORLD2CAM when inverted) :return: Modified tarfile data with poses as ABSOLUTE CAM2WORLD ''' if 'pose' not in tarfile_data: print('[yellow]Warning: No camera poses found in tarfile_data to invert') return tarfile_data inverted_poses = {} for key, pose in tarfile_data['pose'].items(): # pose is (B, 4, 4) tensor B = pose.shape[0] inverted_batch = [] for b in range(B): pose_4x4 = pose[b] # (4, 4) # Extract rotation and translation R = pose_4x4[:3, :3] # (3, 3) t = pose_4x4[:3, 3] # (3,) # Compute inverse: R^T and -R^T @ t R_inv = R.T t_inv = -R_inv @ t # Construct inverted 4x4 matrix pose_inv = torch.zeros_like(pose_4x4) pose_inv[:3, :3] = R_inv pose_inv[:3, 3] = t_inv pose_inv[3, 3] = 1.0 inverted_batch.append(pose_inv) # Stack back to (B, 4, 4) inverted_poses[key] = torch.stack(inverted_batch, dim=0) # Replace poses with inverted versions tarfile_data['pose'] = inverted_poses print('[green]Converted tarfile poses from ABSOLUTE WORLD2CAM to ABSOLUTE CAM2WORLD') return tarfile_data def resize_tarfile_rgb(tarfile_data, target_width, target_height): ''' Resize RGB frames in tarfile data (before converting to VidarDataset format). This operates on raw tarfile data, similar to how invert_tarfile_poses works. :param tarfile_data (dict): Tarfile data dict with 'rgb' key containing (t, c) -> (B, C, H, W) tensors :param target_width (int): Target width :param target_height (int): Target height :return: Modified tarfile data with resized RGB frames ''' import torch.nn.functional as F if 'rgb' not in tarfile_data: print('[yellow]Warning: No RGB data found in tarfile_data to resize') return tarfile_data rgb_dict = tarfile_data['rgb'] resized_rgb = {} # Get original resolution from first frame first_key = next(iter(rgb_dict.keys())) orig_shape = rgb_dict[first_key].shape orig_height, orig_width = orig_shape[2], orig_shape[3] print(f'[cyan]Resizing {len(rgb_dict)} RGB frames from {orig_width}x{orig_height} to {target_width}x{target_height}') for key, rgb_tensor in rgb_dict.items(): # rgb_tensor is (B, C, H, W) if rgb_tensor.shape[2] != target_height or rgb_tensor.shape[3] != target_width: rgb_resized = F.interpolate( rgb_tensor, size=(target_height, target_width), mode='bilinear', align_corners=False ) resized_rgb[key] = rgb_resized else: resized_rgb[key] = rgb_tensor # Replace RGB with resized versions tarfile_data['rgb'] = resized_rgb print(f'[green]Resized all tarfile RGB frames to {target_width}x{target_height}') return tarfile_data # def resize_batch_resolution(input_rgb, base_data_batch, target_width, target_height): # ''' # Resize input RGB and all RGB entries in base_data_batch to target resolution. # Used to override tar/yaml/reference dataset resolution with UI-selected resolution. # :param input_rgb (tensor): Input RGB tensor of shape (T, C, H, W), or None # :param base_data_batch (dict): Data batch containing 'anydata' with 'rgb' dictionary # :param target_width (int): Target width # :param target_height (int): Target height # :return (input_rgb, base_data_batch): Modified tensors/batch with new resolution # ''' # import torch # import torch.nn.functional as F # from einops import rearrange # # Resize input_rgb if present # if input_rgb is not None: # current_shape = input_rgb.shape # print(f'[cyan]Current input shape: {current_shape}') # # Resize to target resolution (T, C, H, W) # input_rgb_resized = F.interpolate( # rearrange(input_rgb, 'T C H W -> 1 (T C) H W'), # size=(target_height, target_width), # mode='bilinear', # align_corners=False # ) # input_rgb = rearrange(input_rgb_resized, '1 (T C) H W -> T C H W', T=current_shape[0], C=current_shape[1]) # print(f'[green]Resized input_rgb to: {input_rgb.shape}') # # Resize RGB data in base_data_batch if present # if 'anydata' in base_data_batch and 'rgb' in base_data_batch['anydata']: # rgb_dict = base_data_batch['anydata']['rgb'] # print(f'[cyan]Resizing {len(rgb_dict)} RGB entries in data batch...') # for key, rgb_tensor in rgb_dict.items(): # # rgb_tensor is (1, C, H, W) # if rgb_tensor.shape[2] != target_height or rgb_tensor.shape[3] != target_width: # rgb_resized = F.interpolate( # rgb_tensor, # size=(target_height, target_width), # mode='bilinear', # align_corners=False # ) # base_data_batch['anydata']['rgb'][key] = rgb_resized # print(f'[green]Resized all RGB entries to {target_width}x{target_height}') # return input_rgb, base_data_batch def pose_to_position_and_direction(pose_matrix): ''' Convert 4x4 pose matrix to camera position and viewing direction. :param pose_matrix (4, 4) tensor or array: Camera extrinsic matrix :return (position, direction, up): Each is (3,) array ''' if isinstance(pose_matrix, torch.Tensor): pose_matrix = pose_matrix.cpu().numpy() # Extract rotation and translation R = pose_matrix[:3, :3] t = pose_matrix[:3, 3] # Camera position in world coordinates # If pose is world-to-camera (extrinsic), camera position is -R^T @ t # If pose is camera-to-world, camera position is just t # Let's try both and use the one that makes more sense position = t # Camera viewing direction (negative Z axis in camera frame, transformed to world) direction = -R[:, 2] # Third column of rotation matrix # Camera up direction (Y axis in camera frame, transformed to world) up = R[:, 1] return position, direction, up def create_camera_frustum_points(position, direction, up, frustum_size=0.3, fov_deg=50.0): ''' Create 3D points for camera frustum visualization. :param position (3,) array: Camera position :param direction (3,) array: Camera viewing direction (unit vector) :param up (3,) array: Camera up direction (unit vector) :param frustum_size (float): Size of frustum :param fov_deg (float): Field of view in degrees :return (5, 3) array: [apex, corner1, corner2, corner3, corner4] ''' # Normalize direction and up direction = direction / (np.linalg.norm(direction) + 1e-8) up = up / (np.linalg.norm(up) + 1e-8) # Camera right direction (cross product) right = np.cross(direction, up) right = right / (np.linalg.norm(right) + 1e-8) # Recompute up to ensure orthogonality up = np.cross(right, direction) up = up / (np.linalg.norm(up) + 1e-8) # Frustum apex at camera position apex = position # Compute frustum corners fov_rad = np.deg2rad(fov_deg) half_height = frustum_size * np.tan(fov_rad / 2.0) half_width = half_height # Four corners of frustum at distance frustum_size center = apex + direction * frustum_size corner1 = center + up * half_height + right * half_width corner2 = center + up * half_height - right * half_width corner3 = center - up * half_height - right * half_width corner4 = center - up * half_height + right * half_width return np.array([apex, corner1, corner2, corner3, corner4]) def visualize_cameras_3d(input_poses_all, output_poses_all, input_label='Input', output_label='Output'): ''' Create 3D plotly visualization of camera poses with XYZ axes. Shows all poses over time but only labels first and last. :param input_poses_all (list): List of all poses over time for input camera (4x4 matrices) :param output_poses_all (list): List of all poses over time for output camera (4x4 matrices) :param input_label (str): Label for input cameras :param output_label (str): Label for output cameras :return: Plotly figure object ''' fig = go.Figure() # Gradient color constants # Trajectory A (input): darker blue → cyan → green (darker for visibility on white) GRAD_A_START = (0.08, 0.20, 0.70, 1.0) # Darker blue GRAD_A_END = (0.03, 0.65, 0.25, 1.0) # Darker green # Trajectory B (output): darker yellow → orange → red (darker for visibility on white) GRAD_B_START = (0.76, 0.64, 0.12, 1.0) # Darker yellow GRAD_B_END = (0.85, 0.15, 0.08, 1.0) # Darker red def interpolate_color(start_rgba, end_rgba, t): '''Interpolate between two RGBA colors. t in [0, 1]''' r = start_rgba[0] + t * (end_rgba[0] - start_rgba[0]) g = start_rgba[1] + t * (end_rgba[1] - start_rgba[1]) b = start_rgba[2] + t * (end_rgba[2] - start_rgba[2]) a = start_rgba[3] + t * (end_rgba[3] - start_rgba[3]) return f'rgba({int(r*255)}, {int(g*255)}, {int(b*255)}, {a})' # Validate that we have poses if not input_poses_all or not output_poses_all: raise ValueError("Input and output pose lists cannot be empty") if input_poses_all[0] is None or input_poses_all[-1] is None: raise ValueError(f"Missing {input_label} poses") if output_poses_all[0] is None or output_poses_all[-1] is None: raise ValueError(f"Missing {output_label} poses") # Convert all poses to numpy and extract positions def extract_positions_and_poses(poses_list): valid_poses = [] positions = [] for pose in poses_list: if pose is not None: if isinstance(pose, torch.Tensor): pose = pose.cpu().numpy() valid_poses.append(pose) positions.append(pose[:3, 3]) return valid_poses, np.array(positions) input_poses, input_positions = extract_positions_and_poses(input_poses_all) output_poses, output_positions = extract_positions_and_poses(output_poses_all) # Check if cameras are fixed (all poses identical over time) def is_camera_fixed(poses, tolerance=1e-4): '''Check if all poses are the same (camera doesn't move)''' if len(poses) <= 1: return True first_pose = poses[0] for pose in poses[1:]: if not np.allclose(pose, first_pose, atol=tolerance): return False return True input_is_fixed = is_camera_fixed(input_poses) output_is_fixed = is_camera_fixed(output_poses) # Compute scene bounds from all positions all_positions = np.vstack([input_positions, output_positions]) center = np.mean(all_positions, axis=0) scene_scale = np.max(np.abs(all_positions - center)) * 1.4 scene_scale = min(max(scene_scale, 1e-2), 1e5) axis_length = scene_scale * 0.22 # Helper function to draw camera axes def draw_camera(pose, color, label=None, show_axes=True, is_endpoint=False): position = pose[:3, 3] R = pose[:3, :3] # Adjust sizes based on whether this is start/end or intermediate if is_endpoint: axis_scale = 1.5 # Bigger for start/end marker_size = 8 line_width_xy = 5 line_width_z = 9 else: axis_scale = 0.6 # Smaller for intermediate marker_size = 3 line_width_xy = 2 line_width_z = 4 if show_axes: # Camera axes in world coordinates scaled_length = axis_length * axis_scale x_axis = R[:, 0] * scaled_length # Right (red) y_axis = R[:, 1] * scaled_length # Up (green) z_axis = R[:, 2] * scaled_length # Forward (blue) # Draw X axis (red) fig.add_trace(go.Scatter3d( x=[position[0], position[0] + x_axis[0]], y=[position[1], position[1] + x_axis[1]], z=[position[2], position[2] + x_axis[2]], mode='lines', line=dict(color='red', width=line_width_xy), showlegend=False, hoverinfo='skip' )) # Draw Y axis (green) fig.add_trace(go.Scatter3d( x=[position[0], position[0] + y_axis[0]], y=[position[1], position[1] + y_axis[1]], z=[position[2], position[2] + y_axis[2]], mode='lines', line=dict(color='green', width=line_width_xy), showlegend=False, hoverinfo='skip' )) # Draw Z axis (viewing direction, blue) fig.add_trace(go.Scatter3d( x=[position[0], position[0] + z_axis[0]], y=[position[1], position[1] + z_axis[1]], z=[position[2], position[2] + z_axis[2]], mode='lines', line=dict(color='blue', width=line_width_z), showlegend=False, hoverinfo='skip' )) # Add camera position marker fig.add_trace(go.Scatter3d( x=[position[0]], y=[position[1]], z=[position[2]], mode='markers', marker=dict(color=color, size=marker_size, symbol='diamond'), showlegend=False, hoverinfo='skip' )) # Add label if provided if label: fig.add_trace(go.Scatter3d( x=[position[0]], y=[position[1]], z=[position[2] + scene_scale * 0.12], mode='text', text=[label], textposition='top center', textfont=dict(color=color, size=10), showlegend=False, hoverinfo='skip' )) # Draw input trajectory if input_is_fixed: # Camera is fixed, only draw one mid_input_color = interpolate_color(GRAD_A_START, GRAD_A_END, 0.5) draw_camera(input_poses[0], mid_input_color, label=f'{input_label} (fixed)', show_axes=True, is_endpoint=True) else: # Draw trajectory line connecting all positions (use middle gradient color) mid_input_color = interpolate_color(GRAD_A_START, GRAD_A_END, 0.5) fig.add_trace(go.Scatter3d( x=input_positions[:, 0], y=input_positions[:, 1], z=input_positions[:, 2], mode='lines', line=dict(color=mid_input_color, width=2), showlegend=False, hoverinfo='skip' )) # Draw all input poses with gradient colors n_input = len(input_poses) for i, pose in enumerate(input_poses): is_first = (i == 0) is_last = (i == n_input - 1) is_endpoint = is_first or is_last show_axes = is_endpoint # Compute gradient color t = i / max(1, n_input - 1) # Normalize to [0, 1] color = interpolate_color(GRAD_A_START, GRAD_A_END, t) label = None if is_first: label = f'{input_label} (first)' elif is_last: label = f'{input_label} (last)' draw_camera(pose, color, label=label, show_axes=show_axes, is_endpoint=is_endpoint) # Draw output trajectory if output_is_fixed: # Camera is fixed, only draw one mid_output_color = interpolate_color(GRAD_B_START, GRAD_B_END, 0.5) draw_camera(output_poses[0], mid_output_color, label=f'{output_label} (fixed)', show_axes=True, is_endpoint=True) else: # Draw trajectory line connecting all positions (use middle gradient color) mid_output_color = interpolate_color(GRAD_B_START, GRAD_B_END, 0.5) fig.add_trace(go.Scatter3d( x=output_positions[:, 0], y=output_positions[:, 1], z=output_positions[:, 2], mode='lines', line=dict(color=mid_output_color, width=2), showlegend=False, hoverinfo='skip' )) # Draw all output poses with gradient colors n_output = len(output_poses) for i, pose in enumerate(output_poses): is_first = (i == 0) is_last = (i == n_output - 1) is_endpoint = is_first or is_last show_axes = is_endpoint # Compute gradient color t = i / max(1, n_output - 1) # Normalize to [0, 1] color = interpolate_color(GRAD_B_START, GRAD_B_END, t) label = None if is_first: label = f'{output_label} (first)' elif is_last: label = f'{output_label} (last)' draw_camera(pose, color, label=label, show_axes=show_axes, is_endpoint=is_endpoint) # Add scene center/origin marker at (0, 0, 0) fig.add_trace(go.Scatter3d( x=[0], y=[0], z=[0], mode='markers', marker=dict(color='black', size=8, symbol='diamond'), name='Scene Center', showlegend=True, hoverinfo='text', hovertext='Origin (0, 0, 0)' )) # Configure layout fig.update_layout( height=600, autosize=True, hovermode='closest', margin=dict(l=0, r=0, b=0, t=30), showlegend=True, legend=dict( yanchor='top', y=0.99, xanchor='left', x=0.01, bgcolor='rgba(255, 255, 255, 0.8)' ), scene=dict( aspectmode='cube', xaxis=dict( title='X', range=[center[0] - scene_scale, center[0] + scene_scale], showgrid=True, gridcolor='lightgray' ), yaxis=dict( title='Y', range=[center[1] - scene_scale, center[1] + scene_scale], showgrid=True, gridcolor='lightgray' ), zaxis=dict( title='Z', range=[center[2] - scene_scale, center[2] + scene_scale], showgrid=True, gridcolor='lightgray' ), camera=dict( eye=dict(x=1.5, y=1.5, z=1.2), center=dict(x=0, y=0, z=0), up=dict(x=0, y=0, z=1) ) ) ) return fig def load_rgb_from_reference_dataset(base_data_batch, cam_idx=1, extract_frames_fn=None): ''' Extract RGB frames from reference dataset (base_data_batch). :param base_data_batch: Data batch containing vidar dict with RGB frames :param cam_idx: Camera index to extract (default 1 for input/context camera) :param extract_frames_fn: Optional function to extract frames (signature: fn(rgb_dict, cam_idx, max_frames)) :return: Tuple of (input_rgb tensor or None, metadata dict, error_msg or None) metadata contains: 'cam_in_count', 'cam_out_count', 'time_indices' ''' if base_data_batch is None: return None, {}, 'Reference dataset not loaded' print(f'[cyan]Loading reference dataset RGB (camera {cam_idx})') rgb_dict = base_data_batch['anydata']['rgb'] # If custom extraction function provided, use it if extract_frames_fn is not None: input_rgb, time_indices, error_msg = extract_frames_fn(rgb_dict, cam_idx=cam_idx, max_frames=None) if input_rgb is None: return None, {}, f'{error_msg} in reference dataset' print(f'[green]Loaded {input_rgb.shape[0]} frames from reference dataset (camera {cam_idx}, frames {time_indices[0]}-{time_indices[-1]})') else: # Default extraction: stack all frames for camera cam_keys = sorted([k for k in rgb_dict.keys() if k[1] == cam_idx]) if len(cam_keys) == 0: return None, {}, f'No camera {cam_idx} frames found in reference dataset' input_rgb = torch.stack([rgb_dict[k][0] for k in cam_keys], dim=0) time_indices = [k[0] for k in cam_keys] print(f'[green]Extracted input_rgb from reference dataset camera {cam_idx}: {input_rgb.shape}') # Collect metadata cam_out_count = len([k for k in rgb_dict.keys() if k[1] == 0]) cam_in_count = len([k for k in rgb_dict.keys() if k[1] == 1]) metadata = { 'cam_in_count': cam_in_count, 'cam_out_count': cam_out_count, 'time_indices': time_indices } return input_rgb, metadata, None def load_rgb_from_yaml_dataset(yaml_path, cam_idx=1, extract_frames_fn=None): ''' Load YAML dataset and extract RGB frames from a random sample. :param yaml_path: Path to YAML dataset config file :param cam_idx: Camera index to extract (default 1 for input/context camera) :param extract_frames_fn: Optional function to extract frames (signature: fn(rgb_dict, cam_idx, max_frames)) :return: Tuple of (input_rgb tensor or None, sample_batch or None, metadata dict, error_msg or None) metadata contains: 'cam_in_count', 'cam_out_count', 'time_indices', 'dataset_size' ''' from custom.dataset.vidar_dataset import VidarDataset from torch.utils.data import DataLoader print(f'[cyan]Loading dataset from YAML: {yaml_path}') try: dataset = VidarDataset( dataset_dir=yaml_path, phase='test', ) print(f'[green]Dataset loaded: {len(dataset)} samples') # Create dataloader with shuffle for random sample dataloader = DataLoader( dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=True, ) # Get random sample sample = next(iter(dataloader)) rgb_dict = sample['anydata']['rgb'] # If custom extraction function provided, use it if extract_frames_fn is not None: input_rgb, time_indices, error_msg = extract_frames_fn(rgb_dict, cam_idx=cam_idx, max_frames=None) if input_rgb is None: return None, None, {}, f'{error_msg} in YAML dataset' print(f'[green]Loaded {input_rgb.shape[0]} frames from YAML dataset (camera {cam_idx}, frames {time_indices[0]}-{time_indices[-1]})') else: # Default extraction: stack all frames for camera cam_keys = sorted([k for k in rgb_dict.keys() if k[1] == cam_idx]) if len(cam_keys) == 0: return None, None, {}, f'No camera {cam_idx} frames found in YAML dataset' input_rgb = torch.stack([rgb_dict[k][0] for k in cam_keys], dim=0) time_indices = [k[0] for k in cam_keys] print(f'[green]Extracted input_rgb from YAML dataset camera {cam_idx}: {input_rgb.shape}') # Collect metadata cam_out_count = len([k for k in rgb_dict.keys() if k[1] == 0]) cam_in_count = len([k for k in rgb_dict.keys() if k[1] == 1]) metadata = { 'cam_in_count': cam_in_count, 'cam_out_count': cam_out_count, 'time_indices': time_indices, 'dataset_size': len(dataset) } return input_rgb, sample, metadata, None except Exception as e: import traceback print(f'[red]Error loading YAML dataset: {e}') print(f'[red]{traceback.format_exc()}') return None, None, {}, f'Failed to load YAML dataset: {str(e)}' def save_camera_poses_to_npy(base_data_batch, output_filepath): ''' Save camera poses (input + output, absolute cam2world) to a .npy file. Uses VIDAR convention: camera 0 = output/target, camera 1 = input/context :param base_data_batch (dict): Data batch containing vidar dict with poses :param output_filepath (str): Full path to output .npy file :return: True if successful, False otherwise ''' try: # Extract poses from base_data_batch and convert to absolute cam2world format if 'anydata' not in base_data_batch or 'pose' not in base_data_batch['anydata']: print('[yellow]Warning: No camera poses found in data batch') return False vidar_poses = base_data_batch['anydata']['pose'] absolute_poses = convert_vidar_to_absolute_format(vidar_poses) # Extract input (cam 1) and output (cam 0) poses over all timesteps all_timesteps = sorted(set([key[0] for key in vidar_poses.keys()])) input_poses = [] # Camera 1 (input/context) output_poses = [] # Camera 0 (output/target) for t in all_timesteps: # Input camera (cam 1) if (t, 1) in absolute_poses: input_poses.append(absolute_poses[(t, 1)][0].cpu().numpy()) # (4, 4) # Output camera (cam 0) if (t, 0) in absolute_poses: output_poses.append(absolute_poses[(t, 0)][0].cpu().numpy()) # (4, 4) # Stack into arrays: (T, 4, 4) input_poses_array = np.stack(input_poses, axis=0) if len(input_poses) > 0 else np.array([]) output_poses_array = np.stack(output_poses, axis=0) if len(output_poses) > 0 else np.array([]) # Save to .npy file poses_dict = { 'input_poses': input_poses_array, # (T, 4, 4) - absolute cam2world 'output_poses': output_poses_array, # (T, 4, 4) - absolute cam2world 'timesteps': all_timesteps, 'format': 'absolute_cam2world', 'convention': 'cam0=output, cam1=input' } np.save(output_filepath, poses_dict) print(f'[green]Saved camera poses (input + output, absolute cam2world): {output_filepath}') print(f'[cyan] Input poses shape: {input_poses_array.shape}, Output poses shape: {output_poses_array.shape}') return True except Exception as e: print(f'[yellow]Warning: Failed to save camera poses: {e}') import traceback print(traceback.format_exc()) return False def extract_and_visualize_cameras(base_data_batch, camera_info_desc=''): ''' Extract camera poses from data batch and create visualization. Uses VIDAR convention: camera 0 = output/target, camera 1 = input/context :param base_data_batch (dict): Data batch containing data dict with poses :param camera_info_desc (str): Description of camera configuration to include in output :return: Tuple of (Plotly figure object or None, extrinsics text string) ''' # Hardcoded camera indices cam_in = 1 # Input/context camera cam_out = 0 # Output/target camera try: pose_dict = base_data_batch.get('anydata', base_data_batch.get('vidar', {})) # Convert all poses from relative world2cam to absolute cam2world if 'pose' not in pose_dict or len(pose_dict['pose']) == 0: return None, 'No camera poses available in data batch' absolute_poses = convert_vidar_to_absolute_format(pose_dict['pose']) # Find all timesteps and sort them all_timesteps = sorted(set([key[0] for key in pose_dict['pose'].keys()])) max_t = all_timesteps[-1] if all_timesteps else 0 # Extract ALL input camera poses over time (using cam_in) input_poses_all = [] for t in all_timesteps: pose = absolute_poses.get((t, cam_in)) if pose is not None and pose.shape[0] > 0: input_poses_all.append(pose[0]) # Take first batch element else: input_poses_all.append(None) # Extract ALL output camera poses over time (using cam_out) output_poses_all = [] for t in all_timesteps: pose = absolute_poses.get((t, cam_out)) if pose is not None and pose.shape[0] > 0: output_poses_all.append(pose[0]) # Take first batch element else: output_poses_all.append(None) # Get first and last for text display input_first_pose = input_poses_all[0] if len(input_poses_all) > 0 else None input_last_pose = input_poses_all[-1] if len(input_poses_all) > 0 else None output_first_pose = output_poses_all[0] if len(output_poses_all) > 0 else None output_last_pose = output_poses_all[-1] if len(output_poses_all) > 0 else None # Build extrinsics text for display extrinsics_lines = [] # Print poses for debugging print('[cyan]=' * 40) print('[cyan]Camera extrinsics:') if input_first_pose is not None: pos, dir, up = pose_to_position_and_direction(input_first_pose) line = f'Input (first) - Pos: [{pos[0]:.2f}, {pos[1]:.2f}, {pos[2]:.2f}], Dir: [{dir[0]:.2f}, {dir[1]:.2f}, {dir[2]:.2f}]' print(f'[blue]{line}') extrinsics_lines.append(line) if input_last_pose is not None: pos, dir, up = pose_to_position_and_direction(input_last_pose) line = f'Input (last) - Pos: [{pos[0]:.2f}, {pos[1]:.2f}, {pos[2]:.2f}], Dir: [{dir[0]:.2f}, {dir[1]:.2f}, {dir[2]:.2f}]' print(f'[cyan]{line}') extrinsics_lines.append(line) if output_first_pose is not None: pos, dir, up = pose_to_position_and_direction(output_first_pose) line = f'Output (first) - Pos: [{pos[0]:.2f}, {pos[1]:.2f}, {pos[2]:.2f}], Dir: [{dir[0]:.2f}, {dir[1]:.2f}, {dir[2]:.2f}]' print(f'[red]{line}') extrinsics_lines.append(line) if output_last_pose is not None: pos, dir, up = pose_to_position_and_direction(output_last_pose) line = f'Output (last) - Pos: [{pos[0]:.2f}, {pos[1]:.2f}, {pos[2]:.2f}], Dir: [{dir[0]:.2f}, {dir[1]:.2f}, {dir[2]:.2f}]' print(f'[yellow]{line}') extrinsics_lines.append(line) print('[cyan]=' * 40) # Build extrinsics text with camera configuration info text_parts = [] if camera_info_desc: text_parts.append(f'📷 {camera_info_desc}') text_parts.append('') # Empty line for spacing text_parts.extend(extrinsics_lines) extrinsics_text = '\n\n'.join(text_parts) if text_parts else 'No camera poses available' # Create visualization with all poses fig = visualize_cameras_3d(input_poses_all, output_poses_all) return fig, extrinsics_text except Exception as e: print(f'[red]Error extracting/visualizing cameras: {e}') import traceback tb_str = traceback.format_exc() print(tb_str) return None, f'❌ Error extracting camera poses:\n{str(e)}\n\nStacktrace:\n{tb_str}'