# Copyright 2026 Toyota Research Institute. All rights reserved. """ Visualization script with minimal dependencies to load a single LBM episode in spartan format and visualize the arm's end-effector poses + action trajectories projected onto RGB images, using camera intrinsics and extrinsics and generate a video. Usage: Get LBM spartan data from s3 (Example): AWS_REGION=us-east-1 aws s3 cp s3://robotics-manip-lbm/efs/data/tasks/BimanualHangMugsOnMugHolderFromDryingRack/riverway/sim/bc/teleop/2024-12-16T11-49-42-05-00/diffusion_spartan/episode_0/processed/ . --recursive Then run: uv run python3 anydata/visualize/vis_spartan.py \ --episode_path data/cv_downloaded/spartan_tiny/efs/data/tasks/BimanualPutRedBellPepperInBin/riverway/sim/bc/teleop/2025-01-06T08-58-31-05-00/diffusion_spartan/episode_131/processed \ --output_dir ./validation_videos \ --camera scene_right_0 \ --max_frames 200 """ import os import argparse import cv2 import numpy as np import torch from tqdm import tqdm from anydata.geometry.camera_utils import points_to_camera_frame, project_to_image from anydata.utils.read import read_yaml from anydata.utils.viz import draw_points_on_image from anydata.utils.write import write_video def visualize_lbm_episode(episode_path, output_dir, max_frames=None, camera='scene_right_0'): """ Load a single LBM episode and visualize end-effector poses on RGB images. Args: episode_path: Path to the processed episode directory output_dir: Directory to save visualization video max_frames: Maximum number of frames to process (None = all frames) """ os.makedirs(output_dir, exist_ok=True) print(f"Loading episode from: {episode_path}") # Load metadata meta_data_file = os.path.join(episode_path, "metadata.yaml") if not os.path.isfile(meta_data_file): raise FileNotFoundError(f"Metadata file not found: {meta_data_file}") meta_data = read_yaml(meta_data_file) print(f"Metadata loaded: {list(meta_data.keys())}") # Load observations observations_file = os.path.join(episode_path, "observations.npz") if not os.path.isfile(observations_file): raise FileNotFoundError(f"Observations file not found: {observations_file}") observations = np.load(observations_file) print(f"Observation keys: {list(observations.keys())}") # Get camera data camera_names = {val: key for key, val in meta_data["camera_id_to_semantic_name"].items()} camera_name = camera if camera_name not in camera_names: raise ValueError(f"Camera {camera_name} not found. Available cameras: {list(camera_names.keys())}") camera_id = camera_names[camera_name] print(f"Using camera: {camera_name} (ID: {camera_id})") # Load camera images and parameters front_img = observations[camera_id] # (N, H, W, 3) intrinsics = np.load(os.path.join(episode_path, "intrinsics.npz"))[camera_id] # (3, 3) extrinsics = np.load(os.path.join(episode_path, "extrinsics.npz"))[camera_id][0] # (4, 4) print(f"Images shape: {front_img.shape}") print(f"Intrinsics shape: {intrinsics.shape}") print(f"Extrinsics shape: {extrinsics.shape}") # Load end-effector poses pose_xyz_left = observations["robot__actual__poses__left::panda__xyz"] # (N, 3) pose_xyz_right = observations["robot__actual__poses__right::panda__xyz"] # (N, 3) robot_position_action_left = observations["robot__desired__poses__left::panda__xyz"] # (N, 3) robot_position_action_right = observations["robot__desired__poses__right::panda__xyz"] print(f"Left EE pose shape: {pose_xyz_left.shape}") print(f"Right EE pose shape: {pose_xyz_right.shape}") print(f"Left EE action shape: {robot_position_action_left.shape}") print(f"Right EE action shape: {robot_position_action_right.shape}") # Load desired joint positions and actions for chunking robot_joint_action_left = observations["robot__desired__joint_position__left::panda"] # (N, 7) robot_joint_action_right = observations["robot__desired__joint_position__right::panda"] # (N, 7) robot_gripper_action_left = observations["robot__desired__grippers__left::panda_hand"] # (N, 1) robot_gripper_action_right = observations["robot__desired__grippers__right::panda_hand"] # (N, 1) # Combine actions for chunking (same format as processing script) position_actions = np.hstack([robot_position_action_left, robot_position_action_right]) # (N, 6) joint_actions = np.hstack([robot_joint_action_left, robot_joint_action_right]) # (N, 14) gripper_actions = np.hstack([robot_gripper_action_left, robot_gripper_action_right]) # (N, 2) actions = np.hstack([position_actions, joint_actions, gripper_actions]) # (N, 22) print(f"Combined actions shape: {actions.shape}") # Process actions with chunking horizon_seconds = 4.0 N_actions = actions.shape[0] chunk_size = int(N_actions / horizon_seconds) ac_dim = actions.shape[1] print(f"Chunk size: {chunk_size}, Action dimension: {ac_dim}") ac_chunks = [] for i in range(0, N_actions): if i + chunk_size > N_actions: # Not enough data to create another chunk, tile last action ac_chunk = np.zeros((1, chunk_size, ac_dim)) ac_chunk[:, :N_actions - i] = actions[i : N_actions].reshape(1, -1, ac_dim) ac_chunk[:, N_actions - i :] = np.tile( actions[N_actions - 1].reshape(1, 1, ac_dim), (1, chunk_size - (N_actions - i), 1) ) else: ac_chunk = actions[i : i + chunk_size].reshape(1, chunk_size, ac_dim) ac_chunks.append(ac_chunk) ac_chunks = np.concatenate(ac_chunks, axis=0) ac_chunks = ac_chunks.astype(np.float32) ac_chunks = np.nan_to_num(ac_chunks, nan=0.0, posinf=0.0, neginf=0.0) print(f"Chunked actions shape: {ac_chunks.shape}") # Extract XYZ actions from chunks left_xyz_act = ac_chunks[:, :, :3] # (N, chunk_size, 3) right_xyz_act = ac_chunks[:, :, 3:6] # (N, chunk_size, 3) # Transform action chunks to camera frame left_xyz_act_flat = left_xyz_act.reshape(-1, 3) # (N*chunk_size, 3) right_xyz_act_flat = right_xyz_act.reshape(-1, 3) # (N*chunk_size, 3) left_xyz_act_cam_flat = points_to_camera_frame(left_xyz_act_flat, extrinsics) right_xyz_act_cam_flat = points_to_camera_frame(right_xyz_act_flat, extrinsics) # Reshape back to (N, chunk_size, 3) left_xyz_act_cam = left_xyz_act_cam_flat.reshape(N_actions, chunk_size, 3) right_xyz_act_cam = right_xyz_act_cam_flat.reshape(N_actions, chunk_size, 3) combined_xyz_act_cam = np.concatenate([left_xyz_act_cam, right_xyz_act_cam], axis=2) # (N, chunk_size, 6) print(f"Combined XYZ action trajectory shape (camera frame): {combined_xyz_act_cam.shape}") # Transform poses to camera frame pose_xyz_left_cam = points_to_camera_frame(pose_xyz_left, extrinsics) # (N, 3) pose_xyz_right_cam = points_to_camera_frame(pose_xyz_right, extrinsics) # (N, 3) # Combine left and right EE poses ee_pose = np.hstack([pose_xyz_left_cam, pose_xyz_right_cam]) # (N, 6) print(f"Combined EE pose in camera frame shape: {ee_pose.shape}") # Determine number of frames to process N = front_img.shape[0] if max_frames is not None: N = min(N, max_frames) print(f"Processing {N} frames...") # Convert 3x3 intrinsics to 3x4 projection matrix if needed intrinsics_proj = intrinsics if intrinsics_proj.shape == (3, 3): intrinsics_proj = np.hstack([intrinsics_proj, np.zeros((3, 1))]) print(f"Intrinsics projection matrix: {intrinsics_proj}") print(f"Extrinsics matrix: {extrinsics}") # Create video frames video_frames = [] for t in tqdm(range(N), desc="Creating visualization"): # Get frame img = front_img[t] # (H, W, 3) if img.max() <= 1.0: # Normalize if needed img = (img * 255).astype(np.uint8) else: img = img.astype(np.uint8) # Get EE pose for this timestep ee_pose_t = ee_pose[t, None, :] # (1, 6) - single pose point # Get EE action for this timestep - use chunked trajectory ee_action_chunk_t = combined_xyz_act_cam[t] # (chunk_size, 6) - full action trajectory # Draw end-effector pose as red dots ee_pose_t_proj = project_to_image(ee_pose_t.reshape(-1, 3), intrinsics_proj) img = draw_points_on_image(img, ee_pose_t_proj, palette="Reds", dot_size=5) # Also display text near the dot (Left_EE and Right_EE) left_ee_px = ee_pose_t_proj[0] right_ee_px = ee_pose_t_proj[1] cv2.putText(img, 'Left_EE', (int(left_ee_px[0]) + 5, int(left_ee_px[1]) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) cv2.putText(img, 'Right_EE', (int(right_ee_px[0]) + 5, int(right_ee_px[1]) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) # Draw end-effector action trajectory as green dots ee_action_chunk_t_proj = project_to_image(ee_action_chunk_t.reshape(-1, 3), intrinsics_proj) img = draw_points_on_image(img, ee_action_chunk_t_proj, palette="Greens", dot_size=3) # Add frame info text on top-left cv2.putText(img, f'Frame {t}/{N}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2) cv2.putText(img, f'Red: EE Pose', (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) cv2.putText(img, f'Green: EE Action Trajectory', (10, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) # Add EE position values as text below left_pos = ee_pose[t, :3] right_pos = ee_pose[t, 3:] cv2.putText(img, f'Left EE: ({left_pos[0]:.2f}, {left_pos[1]:.2f}, {left_pos[2]:.2f})', (10, 400), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) cv2.putText(img, f'Right EE: ({right_pos[0]:.2f}, {right_pos[1]:.2f}, {right_pos[2]:.2f})', (10, 420), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) video_frames.append(img) # Save video if video_frames: video_tensor = torch.stack([torch.from_numpy(frame) for frame in video_frames]) episode_name = os.path.basename(os.path.dirname(episode_path)) video_path = os.path.join(output_dir, f'{episode_name}_{camera_name}_ee_visualization.mp4') print(f"Saving video: {video_path}") write_video(video_path, video_tensor, fps=10) print(f"Video saved successfully!") else: print("No frames to save!") def main(): parser = argparse.ArgumentParser( description='Visualize end-effector poses from a single LBM episode') parser.add_argument('--episode_path', type=str, required=True, help='Path to the processed episode directory') parser.add_argument('--output_dir', type=str, default='./validation_videos', help='Output directory for visualization video') parser.add_argument('--max_frames', type=int, default=None, help='Maximum frames to process (default: all frames)') parser.add_argument('--camera', type=str, default='scene_right_0', help='Camera to visualize (default: scene_right_0)') args = parser.parse_args() # Validate camera against names available in this episode's metadata meta_data = read_yaml(os.path.join(args.episode_path, 'metadata.yaml')) available_cameras = list(meta_data['camera_id_to_semantic_name'].values()) if args.camera not in available_cameras: parser.error(f"--camera '{args.camera}' not found. Available: {available_cameras}") print(f"Episode path: {args.episode_path}") print(f"Output directory: {args.output_dir}") print(f"Camera: {args.camera}") visualize_lbm_episode( args.episode_path, args.output_dir, args.max_frames, args.camera, ) print("Visualization complete!") if __name__ == "__main__": main()