Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| #!/usr/bin/env python3 | |
| """ | |
| Visualize OAK-D RGBD SLAM trajectory using Rerun. | |
| Usage: | |
| uv run python scripts/visualize/visualize_rgbd_trajectory.py -i <episode_dir> | |
| uv run python scripts/visualize/visualize_rgbd_trajectory.py -i <episode_dir> --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]))) | |
| 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() | |