"""Before/after visualization of the wrist-cam GT-keypoint projection bug. Renders 10 frames spread through one AR-view episode. For each frame, projects the future 32 UMI positions into (a) the scene camera and (b) the wrist camera. The wrist projection is done TWICE — once with the buggy composition `umi_pose @ inv(T_wrist_cam_umi)` (my original dataset_arview.py) and once with the correct composition `umi_pose @ T_wrist_cam_umi` (matches viewer's render_dataset_frame.py:456 and the recovery-formula inverse). Output: /data/cameron/repos/smithbot_v4/preprocess/kp_diff_viz/ f{frame:06d}.jpg — a 3-panel horizontal composite per frame (scene | wrist-BUGGY | wrist-FIXED) """ from __future__ import annotations import json from pathlib import Path import cv2 import numpy as np DATASET = Path("/data/cameron/puget_ar_datasets/testing_cube_in_bowl") EPISODES_JSON = Path("/data/cameron/ar_renders/testing_cube_in_bowl/frames/episodes.json") CAL = Path("/data/cameron/claude_feetech_controller/ar_view/wrist_umi_calibration.json") OUT_DIR = Path("/data/cameron/repos/smithbot_v4/preprocess/kp_diff_viz") N_WINDOW = 32 # matches training config N_SAMPLE_FRAMES = 10 # spread evenly through the chosen episode EPISODE_INDEX = 1 # 4 episodes total; picking the 38-frame one def project(xyz_world: np.ndarray, K: np.ndarray, T_w2c: np.ndarray) -> np.ndarray: """xyz_world: (T, 3) → (T, 2) pixel coords (may be off-image).""" homo = np.concatenate([xyz_world, np.ones((len(xyz_world), 1))], axis=1) cam = (T_w2c @ homo.T).T[:, :3] z = cam[:, 2:3] front = z[:, 0] > 1e-3 z = np.where(z <= 1e-3, 1e-3, z) pix = (K @ (cam / z).T).T[:, :2] return pix, front def draw_keypoints(img: np.ndarray, pix: np.ndarray, front: np.ndarray, color_start=(255, 255, 0), color_end=(0, 128, 255)): """Draw a fading polyline through the T keypoints on img. Yellow→orange fades from now to horizon. Points behind camera or off-image get skipped.""" h, w = img.shape[:2] prev = None for i, (x, y) in enumerate(pix): if not front[i] or not (0 <= x < w and 0 <= y < h): prev = None continue alpha = i / max(1, len(pix) - 1) color = tuple(int(color_start[c] * (1 - alpha) + color_end[c] * alpha) for c in range(3)) cv2.circle(img, (int(x), int(y)), 4, color, -1, cv2.LINE_AA) if prev is not None: cv2.line(img, prev, (int(x), int(y)), color, 1, cv2.LINE_AA) prev = (int(x), int(y)) return img def label(img: np.ndarray, txt: str, color=(255, 255, 255)): cv2.rectangle(img, (0, 0), (img.shape[1], 24), (0, 0, 0), -1) cv2.putText(img, txt, (6, 17), cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 1, cv2.LINE_AA) return img def main(): OUT_DIR.mkdir(parents=True, exist_ok=True) meta = json.loads((DATASET / "meta.json").read_text()) K_scene = np.asarray(meta["scene"]["K_save"], dtype=np.float64) K_wrist = np.asarray(meta["wrist"]["K_save"], dtype=np.float64) T_wc_umi = np.asarray(json.loads(CAL.read_text())["T_wrist_cam_umi"], dtype=np.float64) eps = json.loads(EPISODES_JSON.read_text())["episodes"] ep = eps[EPISODE_INDEX] e_start, e_end = int(ep["start"]), int(ep["end"]) print(f"episode {ep['id']} frames [{e_start}..{e_end}] len={e_end-e_start+1}") # Preload all state npzs so we can look up future umi_pose per frame. def _load_state(i): d = np.load(DATASET / "state" / f"{i:06d}.npz") return { "scene_cam_pose": np.asarray(d["scene_cam_pose"], dtype=np.float64), "wrist_cam_pose": np.asarray(d["wrist_cam_pose"], dtype=np.float64), "umi_pose": np.asarray(d["umi_pose"], dtype=np.float64), "umi_visible": bool(d["umi_visible"]), } total_frames = int(meta["num_frames"]) state = [_load_state(i) for i in range(total_frames)] # 10 frames spread through the episode. frame_ids = np.linspace(e_start, e_end - 1, N_SAMPLE_FRAMES).round().astype(int).tolist() for f in frame_ids: # Future 32 UMI positions, clipped to episode end (matches training clip). future_ids = np.clip(np.arange(1, N_WINDOW + 1) + f, 0, e_end) umi_xyz = np.stack([state[i]["umi_pose"][:3, 3] for i in future_ids], axis=0) # Scene camera pose = per-frame (matches viewer; training uses median but # per-frame is fine here since the scene cam is physically fixed anyway). scene_bgr = cv2.imread(str(DATASET / "scene" / f"{f:06d}.jpg")) wrist_bgr = cv2.imread(str(DATASET / "wrist" / f"{f:06d}.jpg")) umi_pose = state[f]["umi_pose"] T_w2c_scene = np.linalg.inv(state[f]["scene_cam_pose"]) T_world_wrist_BUGGY = umi_pose @ np.linalg.inv(T_wc_umi) # OLD (buggy) T_world_wrist_FIXED = umi_pose @ T_wc_umi # NEW (correct) T_w2c_wrist_BUGGY = np.linalg.inv(T_world_wrist_BUGGY) T_w2c_wrist_FIXED = np.linalg.inv(T_world_wrist_FIXED) # Project + draw on each view. p_scene, f_scene = project(umi_xyz, K_scene, T_w2c_scene) p_wB, f_wB = project(umi_xyz, K_wrist, T_w2c_wrist_BUGGY) p_wF, f_wF = project(umi_xyz, K_wrist, T_w2c_wrist_FIXED) panel_scene = draw_keypoints(scene_bgr.copy(), p_scene, f_scene) panel_wrist_B = draw_keypoints(wrist_bgr.copy(), p_wB, f_wB) panel_wrist_F = draw_keypoints(wrist_bgr.copy(), p_wF, f_wF) panel_scene = label(panel_scene, f"scene f={f}", (255, 255, 255)) panel_wrist_B = label(panel_wrist_B, f"wrist BUGGY (umi@inv(T_wc_umi))", (100, 100, 255)) panel_wrist_F = label(panel_wrist_F, f"wrist FIXED (umi@T_wc_umi)", (100, 255, 100)) # Pad panels to same height for hstack. H = max(panel_scene.shape[0], panel_wrist_B.shape[0], panel_wrist_F.shape[0]) def pad(im): if im.shape[0] == H: return im pad_rows = H - im.shape[0] return np.vstack([im, np.zeros((pad_rows, im.shape[1], 3), dtype=np.uint8)]) composite = np.hstack([pad(panel_scene), pad(panel_wrist_B), pad(panel_wrist_F)]) cv2.imwrite(str(OUT_DIR / f"f{f:06d}.jpg"), composite, [cv2.IMWRITE_JPEG_QUALITY, 90]) print(f" f={f}: scene({panel_scene.shape[1]}) + wristB({panel_wrist_B.shape[1]}) + wristF({panel_wrist_F.shape[1]}) → {composite.shape[1]}px") # Also stack all 10 vertically into ONE big image for easy scroll. frames = [cv2.imread(str(OUT_DIR / f"f{f:06d}.jpg")) for f in frame_ids] stack = np.vstack(frames) cv2.imwrite(str(OUT_DIR / "all_frames_stacked.jpg"), stack, [cv2.IMWRITE_JPEG_QUALITY, 90]) print(f"\nwrote {len(frame_ids)} frames + stacked composite → {OUT_DIR}/") print(f" browse: https://omidlab.net/browse/repos/smithbot_v4/preprocess/kp_diff_viz/") if __name__ == "__main__": main()