#!/usr/bin/env python3 """ Visualize OAK-D RGBD SLAM trajectory using Rerun. Usage: uv run python scripts/visualize/visualize_rgbd_trajectory.py -i uv run python scripts/visualize/visualize_rgbd_trajectory.py -i --video-skip 3 """ import json import sys import time from pathlib import Path from grabette_postprocess.episode_files import resolve import click import cv2 import numpy as np import pandas as pd import rerun as rr import rerun.blueprint as rrb from scipy.spatial.transform import Rotation def _gravity_align(positions: np.ndarray, quaternions: np.ndarray, oak_dir: Path, n_samples: int = 50, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Rotate trajectory so gravity (in WORLD frame) points to world -Z. Critical: we have to compute gravity in the *world frame*, not the camera frame, before computing the align rotation. With rtabmap's IMU-init, the first pose's orientation is non-identity (world ≠ camera at t=0), so camera-frame gravity and world-frame gravity differ by that initial rotation. Pipeline: 1. Average first `n_samples` accel readings → gravity in IMU frame. 2. Transform to optical camera frame via calib's imu_to_cam. 3. Rotate through the first pose's orientation to get gravity in the trajectory's world frame. 4. Find shortest-arc rotation R_align that maps world-gravity to (0,0,-1). 5. Apply R_align (left-multiply) to all poses. Returns (positions, quaternions, g_world_unit) where g_world_unit is the measured gravity direction in the original (pre-align) world frame. """ imu_acc = pd.read_csv(oak_dir / "imu_acc.csv") g_imu = imu_acc[["ax", "ay", "az"]].iloc[:n_samples].mean().to_numpy() calib = json.loads((oak_dir / "calib_offline.json").read_text()) R_imu_to_cam = np.array(calib["imu_to_cam"])[:3, :3] # Accelerometer at rest = SPECIFIC FORCE (points UP). Negate for physical # gravity (DOWN) so the align-to-(0,0,-1) below gives a Z-UP world. g_cam = -(R_imu_to_cam @ g_imu) g_cam_unit = g_cam / np.linalg.norm(g_cam) # Project gravity into world frame using the FIRST pose's orientation. R_world_from_cam0 = Rotation.from_quat(quaternions[0]) g_world_unit = R_world_from_cam0.apply(g_cam_unit) # Align world gravity to (0, 0, -1). R_align, _ = Rotation.align_vectors([[0.0, 0.0, -1.0]], [g_world_unit]) new_pos = R_align.apply(positions) new_quats = (R_align * Rotation.from_quat(quaternions)).as_quat() return new_pos, new_quats, g_world_unit def _load_oak_imu(oak_dir: Path) -> dict | None: """Load IMU from OAK imu_acc.csv / imu_gyro.csv (timestamp_ns, x, y, z).""" acc_path = oak_dir / "imu_acc.csv" gyro_path = oak_dir / "imu_gyro.csv" if not acc_path.is_file() or not gyro_path.is_file(): return None def _read(path): samples = [] with open(path) as f: next(f) # skip header for line in f: p = line.strip().split(',') if len(p) < 4: continue samples.append({ 'timestamp': float(p[0]) * 1e-9, 'value': [float(p[1]), float(p[2]), float(p[3])], }) return samples accel = _read(acc_path) gyro = _read(gyro_path) if not accel and not gyro: return None return {'accel': accel, 'gyro': gyro} def _log_oak_imu(imu_data: dict, t0: float): rr.log("imu/accelerometer/x", rr.SeriesLines(colors=[255, 0, 0], names="accel_x"), static=True) rr.log("imu/accelerometer/y", rr.SeriesLines(colors=[0, 255, 0], names="accel_y"), static=True) rr.log("imu/accelerometer/z", rr.SeriesLines(colors=[0, 0, 255], names="accel_z"), static=True) rr.log("imu/gyroscope/x", rr.SeriesLines(colors=[255, 128, 0], names="gyro_x"), static=True) rr.log("imu/gyroscope/y", rr.SeriesLines(colors=[128, 255, 0], names="gyro_y"), static=True) rr.log("imu/gyroscope/z", rr.SeriesLines(colors=[0, 128, 255], names="gyro_z"), static=True) for s in imu_data['accel']: rr.set_time("time", timestamp=s['timestamp'] - t0) v = s['value'] rr.log("imu/accelerometer/x", rr.Scalars(float(v[0]))) rr.log("imu/accelerometer/y", rr.Scalars(float(v[1]))) rr.log("imu/accelerometer/z", rr.Scalars(float(v[2]))) for s in imu_data['gyro']: rr.set_time("time", timestamp=s['timestamp'] - t0) v = s['value'] rr.log("imu/gyroscope/x", rr.Scalars(float(v[0]))) rr.log("imu/gyroscope/y", rr.Scalars(float(v[1]))) rr.log("imu/gyroscope/z", rr.Scalars(float(v[2]))) @click.command() @click.option('-i', '--input_dir', 'episode_dir', required=True, type=click.Path(exists=True), help='Episode directory to visualize') @click.option('--show-video/--no-video', default=True, help='Show camera frames') @click.option('--video-skip', default=5, help='Show every Nth frame') @click.option('--app-id', default='grabette_viz', help='Rerun application ID') @click.option('--gravity-align/--no-gravity-align', default=False, help='Re-align trajectory so first IMU gravity maps to world -Z. ' 'Normally NOT needed — run_oak_slam already does this when ' 'producing camera_trajectory.csv. Use for legacy CSVs that ' 'predate the in-pipeline gravity align.') def main(episode_dir, show_video, video_skip, app_id, gravity_align): """Visualize OAK SLAM trajectory from a processed episode directory.""" episode_dir = Path(episode_dir) oak_dir = episode_dir / "oak" traj_csv = episode_dir / "camera_trajectory.csv" if not traj_csv.exists(): print(f"Error: No camera_trajectory.csv in {episode_dir}") sys.exit(1) # --- Load trajectory --- print(f"Loading trajectory from {traj_csv.name}...") df_all = pd.read_csv(traj_csv) df_valid = df_all[~df_all['is_lost'].astype(bool)].copy() n_total, n_tracked = len(df_all), len(df_valid) print("\n=== SLAM Statistics ===") print(f" Frames: {n_tracked}/{n_total} tracked ({100*n_tracked/n_total:.1f}%)") print(f" Lost: {n_total - n_tracked} frames") if n_tracked > 0: duration = df_valid.iloc[-1]['timestamp'] - df_valid.iloc[0]['timestamp'] print(f" Duration: {duration:.2f}s") positions = df_valid[['x', 'y', 'z']].to_numpy().copy() quaternions = df_valid[['q_x', 'q_y', 'q_z', 'q_w']].to_numpy().copy() if n_tracked == 0: print("Error: No valid poses found!") sys.exit(1) g_unit_for_log = None # gravity direction in trajectory frame (pre-align) if gravity_align: if not oak_dir.is_dir(): print("Error: --gravity-align requires the oak/ subdir (run convert_episode_to_oak.py first)") sys.exit(1) positions, quaternions, g_unit_for_log = _gravity_align( positions, quaternions, oak_dir, ) print("\n=== Gravity Alignment ===") print(f" Measured g in SLAM camera frame (unit): " f"[{g_unit_for_log[0]:+.3f}, {g_unit_for_log[1]:+.3f}, {g_unit_for_log[2]:+.3f}]") print(" Rotated so that direction now maps to [0, 0, -1] (world -Z)") # Also rotate the position/quaternion columns we still need from df_valid df_valid = df_valid.copy() df_valid.loc[:, ['x', 'y', 'z']] = positions df_valid.loc[:, ['q_x', 'q_y', 'q_z', 'q_w']] = quaternions print("\n=== Trajectory Statistics ===") for ax, name in enumerate(['X', 'Y', 'Z']): lo, hi = positions[:, ax].min(), positions[:, ax].max() print(f" {name}: [{lo:.4f}, {hi:.4f}] range={hi-lo:.4f}") distances = np.linalg.norm(np.diff(positions, axis=0), axis=1) print(f" Total path length: {np.sum(distances):.3f} m") print(f" Displacement: {np.linalg.norm(positions[-1] - positions[0]):.3f} m\n") # Normalize timestamps to start from 0 so Rerun timeline is readable t0 = float(df_all['timestamp'].iloc[0]) df_all = df_all.copy() df_all['timestamp'] = df_all['timestamp'] - t0 # --- Load OAK IMU --- imu_data = _load_oak_imu(oak_dir) if oak_dir.is_dir() else None if imu_data: print(f"IMU: {len(imu_data['accel'])} accel, {len(imu_data['gyro'])} gyro samples") else: print("IMU: not found (expected oak/imu_acc.csv and oak/imu_gyro.csv)") # --- Open OAK left video (what the SLAM actually saw) --- # We use dcam_left.mp4 — NOT raw_video.mp4 — because: # 1. raw_video.mp4 is the RPi fisheye camera, not the SLAM input. # 2. dcam_left.mp4 is the rectified SLAM input; its frames are 1:1 # with the trajectory (encoder order == SLAM frame_idx), so no # timestamp arithmetic is needed. video_path = resolve(episode_dir, "dcam_left.mp4") print(f"Looking for video at: {video_path} (exists={video_path.exists()})") video_cap = None video_fps = None if show_video and video_path.exists(): video_cap = cv2.VideoCapture(str(video_path)) if not video_cap.isOpened(): print("Warning: Could not open dcam_left.mp4") video_cap = None else: video_fps = video_cap.get(cv2.CAP_PROP_FPS) n_vframes = int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT)) vid_w = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) vid_h = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) print(f"OAK left video: {n_vframes} frames at {video_fps:.2f} fps ({vid_w}x{vid_h})") elif show_video: print("Video: dcam_left.mp4 not found, camera feed disabled") # --- Initialize Rerun --- rr.init(app_id, spawn=True) time.sleep(0.5) rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True) if imu_data: print("Logging IMU data...") _log_oak_imu(imu_data, t0) # --- Blueprint --- try: top_views = [rrb.Spatial3DView(name="3D View", origin="/world")] if video_cap is not None: top_views.append(rrb.Spatial2DView(name="Camera", origin="/camera_feed")) bottom_views = [] if imu_data: bottom_views.append(rrb.TimeSeriesView(name="Accelerometer", origin="/imu/accelerometer")) bottom_views.append(rrb.TimeSeriesView(name="Gyroscope", origin="/imu/gyroscope")) if bottom_views: blueprint = rrb.Blueprint( rrb.Vertical(rrb.Horizontal(*top_views), rrb.Horizontal(*bottom_views)) ) else: blueprint = rrb.Blueprint(rrb.Horizontal(*top_views)) rr.send_blueprint(blueprint) except Exception as e: print(f"Warning: Could not send blueprint: {e}") # --- Static elements --- axis_len = 0.5 rr.log("world/axes", rr.Arrows3D( origins=[[0, 0, 0]] * 3, vectors=[[axis_len, 0, 0], [0, axis_len, 0], [0, 0, axis_len]], colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]], ), static=True) rr.log("world/trajectory_full", rr.LineStrips3D(positions, colors=[0, 255, 0]), static=True) rr.log("world/start", rr.Points3D(positions[0], colors=[0, 255, 0], radii=0.01), static=True) rr.log("world/end", rr.Points3D(positions[-1], colors=[255, 0, 0], radii=0.01), static=True) # Gravity arrow in WORLD frame: rotate the camera-frame gravity vector # by the camera's first orientation to express it in world coords. After # `--gravity-align`, this should be ~(0, 0, -1) by construction (since the # align step now operates on world-frame gravity, not camera-frame). g_world = None if oak_dir.is_dir(): try: imu_acc = pd.read_csv(oak_dir / "imu_acc.csv") g_imu = imu_acc[["ax", "ay", "az"]].iloc[:50].mean().to_numpy() calib = json.loads((oak_dir / "calib_offline.json").read_text()) R_imu_to_cam = np.array(calib["imu_to_cam"])[:3, :3] # Negate: accel at rest is specific force (UP); physical gravity is DOWN. g_cam = -(R_imu_to_cam @ g_imu) g_cam /= np.linalg.norm(g_cam) R_world_from_cam = Rotation.from_quat(quaternions[0]) g_world = R_world_from_cam.apply(g_cam) except Exception as e: print(f"(gravity arrow disabled: {e})") if g_world is not None: print(f" Gravity in WORLD frame (should be ~(0,0,-1) with --gravity-align): " f"[{g_world[0]:+.3f}, {g_world[1]:+.3f}, {g_world[2]:+.3f}]") rr.log("world/gravity", rr.Arrows3D( origins=[positions[0].tolist()], vectors=[(g_world * 0.3).tolist()], colors=[[255, 255, 0]], labels=["g (world)"], ), static=True) # --- Build pose lookup --- pose_map = {} for _, row in df_valid.iterrows(): pose_map[int(row['frame_idx'])] = { 'pos': np.array([row['x'], row['y'], row['z']]), 'quat': np.array([row['q_x'], row['q_y'], row['q_z'], row['q_w']]), } # --- Animate --- print(f"Visualizing {n_total} frames (skip={video_skip})...") trajectory_so_far = [] cam_axis_len = 0.1 for frame_i in range(n_total): if frame_i % video_skip != 0: continue row = df_all.iloc[frame_i] frame_idx = int(row['frame_idx']) t = float(row['timestamp']) rr.set_time("time", timestamp=t) if frame_idx in pose_map: p = pose_map[frame_idx] pos, quat = p['pos'], p['quat'] rr.log("world/camera", rr.Transform3D(translation=pos.tolist(), quaternion=quat.tolist())) rr.log("world/current_position", rr.Points3D(pos, colors=[255, 0, 0], radii=0.005)) trajectory_so_far.append(pos) if len(trajectory_so_far) > 1: rr.log("world/trajectory_history", rr.LineStrips3D(np.array(trajectory_so_far), colors=[0, 128, 255])) rot = Rotation.from_quat(quat) rr.log("world/camera_axes", rr.Arrows3D( origins=[pos, pos, pos], vectors=[rot.apply([cam_axis_len, 0, 0]), rot.apply([0, cam_axis_len, 0]), rot.apply([0, 0, cam_axis_len])], colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]], )) if video_cap is not None: # dcam_left.mp4 frames are 1:1 with SLAM trajectory frame_idx — # both came from the same StereoDepth output in encoding order. video_cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) ret, frame = video_cap.read() if ret: frame_disp = cv2.resize(frame, (640, 400), interpolation=cv2.INTER_AREA) rr.log("camera_feed", rr.Image(frame_disp)) if frame_i % (video_skip * 10) == 0: n_vframes_logged = frame_i // video_skip + 1 print(f" Frame {frame_i}/{n_total} video_idx={frame_idx if video_cap else '-'} logged={n_vframes_logged}", end='\r') print("\nVisualization complete.") print(" Green line: full trajectory (static)") print(" Blue line: trajectory up to current time") print(" RGB arrows: camera X/Y/Z axes") if video_cap is not None: video_cap.release() print("\nPress Ctrl+C to exit...") try: while True: time.sleep(1) except KeyboardInterrupt: print("\nExiting.") if __name__ == "__main__": main()