""" Extract T-Rex 250-point tracks for LeRobot v2 episodes with SAM2 + CoTracker. Frame-0 masks: one SAM2 predict per hand/view (fixed prompts, no appearance auto-repair). Then CoTracker tracks: - head_left: 100 pts (left hand/arm 50 + right hand/arm 50) - each wrist: 75 pts (fixed 5×5 background grid 25 + hand 50) Tune prompts in ``trex_track/sam2_prompt_hands.py``. Example:: CUDA_VISIBLE_DEVICES=4 python scripts/extract_track.py \\ --dataset-root data/trex_small --episode-index 0 \\ --cotracker-device cuda:0 --sam2-device cuda:0 """ from __future__ import annotations import argparse import os import sys import tempfile from pathlib import Path from typing import NamedTuple import numpy as np _SCRIPT_DIR = Path(__file__).resolve().parent _DREAMZERO_ROOT = _SCRIPT_DIR.parent if str(_SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(_SCRIPT_DIR)) from trex_track.layout import ( # noqa: E402 NUM_COMBINED_POINTS, NUM_HEAD_LEFT, NUM_HEAD_PER_HAND as NUM_HAND_POINTS, NUM_HEAD_POINTS, NUM_WRIST_BACKGROUND as NUM_WRIST_GRID, NUM_WRIST_HAND, NUM_WRIST_POINTS, POINT_SLICES, TRACK_LAYOUT_VERSION, VIEW_ORDER, identity_metadata, ) VIDEO_FOLDERS = { "head_left": "observation.images.head_left", "left_wrist": "observation.images.left_wrist", "right_wrist": "observation.images.right_wrist", } DEFAULT_OPENPI_ROOT = Path("/scratch2/home/zhicao/openpi") DEFAULT_CALIB = _DREAMZERO_ROOT / "assets" / "trex_camera_calib.json" DEFAULT_SAM2_MODEL = os.environ.get("SAM2_MODEL", "facebook/sam2-hiera-large") DEFAULT_SAM2_LIBS = os.environ.get("SAM2_LIBS", "/scratch1/home/zhicao/physctrl/libs") class TrackingRuntime(NamedTuple): """Lazily-created heavy models shared across an episode batch.""" calib: dict out_hw: tuple[int, int] cotracker_model: object cotracker_device: object sam2_predictor: object def _ensure_openpi_on_path(openpi_root: Path) -> Path: root = openpi_root.expanduser().resolve() droid = root / "droid" if not droid.is_dir(): raise FileNotFoundError(f"openpi droid package not found: {droid}") p = str(droid) if p not in sys.path: sys.path.insert(0, p) return droid def _enable_cotracker_sdpa_attention(openpi_root: str | Path) -> None: """Replace CoTracker's quadratic-memory attention with PyTorch SDPA. The original implementation materializes an ``[B,H,T,T]`` attention matrix. Long T-Rex episodes can therefore require more than 80 GB even on an otherwise empty H100. SDPA uses Flash Attention for CUDA BF16 inputs, preserving full-sequence attention without materializing that matrix. """ import torch.nn.functional as F cotracker_root = Path(openpi_root).expanduser().resolve() / "co-tracker" cotracker_path = str(cotracker_root) if not cotracker_root.is_dir(): raise FileNotFoundError(f"CoTracker package not found: {cotracker_root}") if cotracker_path not in sys.path: sys.path.insert(0, cotracker_path) from cotracker.models.core.cotracker.blocks import Attention if bool(getattr(Attention, "_trex_sdpa_enabled", False)): return def _sdpa_forward(self, x, context=None, attn_bias=None): batch, query_steps, _ = x.shape heads = int(self.heads) query = self.to_q(x) inner_dim = int(query.shape[-1]) head_dim = inner_dim // heads query = query.reshape(batch, query_steps, heads, head_dim).transpose(1, 2) context = x if context is None else context key, value = self.to_kv(context).chunk(2, dim=-1) context_steps = int(context.shape[1]) key = key.reshape(batch, context_steps, heads, head_dim).transpose(1, 2) value = value.reshape(batch, context_steps, heads, head_dim).transpose(1, 2) attended = F.scaled_dot_product_attention( query, key, value, attn_mask=attn_bias, dropout_p=0.0, is_causal=False, ) attended = attended.transpose(1, 2).reshape(batch, query_steps, inner_dim) return self.to_out(attended) Attention.forward = _sdpa_forward Attention._trex_sdpa_enabled = True def normalize_tracks_xy(tracks: np.ndarray, img_w: int, img_h: int) -> np.ndarray: out = np.asarray(tracks, dtype=np.float32).copy() w, h = max(float(img_w), 1.0), max(float(img_h), 1.0) out = np.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0) out[..., 0] = np.clip(out[..., 0] / w, 0.0, 1.0) out[..., 1] = np.clip(out[..., 1] / h, 0.0, 1.0) return out.astype(np.float32, copy=False) def normalize_track_result( tracks: np.ndarray, visibility: np.ndarray, img_w: int, img_h: int, ) -> tuple[np.ndarray, np.ndarray]: """Normalize XY and clear visibility for non-finite/out-of-frame points.""" pixels = np.asarray(tracks, dtype=np.float32) vis = np.asarray(visibility, dtype=np.float32) if pixels.ndim != 3 or pixels.shape[-1] != 2: raise ValueError(f"tracks must be (T,N,2), got {pixels.shape}") if vis.shape != pixels.shape[:2]: raise ValueError(f"visibility {vis.shape} does not match tracks {pixels.shape}") finite = np.isfinite(pixels).all(axis=-1) in_frame = ( (pixels[..., 0] >= 0.0) & (pixels[..., 0] < float(img_w)) & (pixels[..., 1] >= 0.0) & (pixels[..., 1] < float(img_h)) ) clean_vis = ((vis > 0.5) & finite & in_frame).astype(np.float32) return normalize_tracks_xy(pixels, img_w, img_h), clean_vis def load_episode_videos( dataset_root: Path, episode_index: int, *, out_hw: tuple[int, int], ) -> dict[str, np.ndarray]: import cv2 chunk = episode_index // 1000 frames_by_view: dict[str, list[np.ndarray]] = {k: [] for k in VIEW_ORDER} for view in VIEW_ORDER: rel = VIDEO_FOLDERS[view] video_path = ( dataset_root / "videos" / f"chunk-{chunk:03d}" / rel / f"episode_{episode_index:06d}.mp4" ) if not video_path.is_file(): raise FileNotFoundError(video_path) cap = cv2.VideoCapture(str(video_path)) while True: ok, bgr = cap.read() if not ok: break rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) if (rgb.shape[0], rgb.shape[1]) != out_hw: rgb = cv2.resize(rgb, (out_hw[1], out_hw[0]), interpolation=cv2.INTER_LINEAR) frames_by_view[view].append(rgb) cap.release() if not frames_by_view[view]: raise RuntimeError(f"empty video: {video_path}") return {v: np.stack(frames_by_view[v], axis=0) for v in VIEW_ORDER} def load_episode_states(dataset_root: Path, episode_index: int) -> tuple[np.ndarray, str]: import pandas as pd chunk = episode_index // 1000 pq_path = dataset_root / "data" / f"chunk-{chunk:03d}" / f"episode_{episode_index:06d}.parquet" df = pd.read_parquet(pq_path) states = np.stack([np.asarray(x, dtype=np.float64) for x in df["observation.state"].values], axis=0) task = "" if "annotation.task" in df.columns: task = str(df["annotation.task"].iloc[0]) return states, task def _run_cotracker_window( model, video_hwc: np.ndarray, query_xy: np.ndarray, device: object, ) -> tuple[np.ndarray, np.ndarray]: """Run one bounded CoTracker window and immediately release its GPU tensors.""" import torch video_np = np.asarray(video_hwc, dtype=np.uint8) queries_np = np.zeros((int(query_xy.shape[0]), 3), dtype=np.float32) queries_np[:, 1:] = np.asarray(query_xy, dtype=np.float32) device_type = torch.device(device).type video_dtype = torch.bfloat16 if device_type == "cuda" else torch.float32 video = ( torch.from_numpy(video_np) .permute(0, 3, 1, 2) .unsqueeze(0) .to(device=device, dtype=video_dtype) ) queries = torch.from_numpy(queries_np).unsqueeze(0).to(device) with torch.inference_mode(), torch.autocast( device_type=device_type, dtype=torch.bfloat16, enabled=device_type == "cuda", ): pred_tracks, pred_vis = model( video, queries=queries, backward_tracking=False, ) tracks = pred_tracks[0].detach().cpu().numpy().astype(np.float32) visibility = pred_vis[0].detach().cpu().numpy() visibility = (visibility > 0.5).astype(np.float32) del video, queries, pred_tracks, pred_vis if device_type == "cuda": torch.cuda.empty_cache() return tracks, visibility def _run_cotracker(model, video_hwc: np.ndarray, query_xy: np.ndarray, device: object): """Track frame-0 queries in bounded, overlapping temporal windows. ``CoTrackerPredictor`` only copies backward predictions into frames before each query timestamp. Every query here starts at frame zero, so ``backward_tracking=True`` cannot change the result and nearly doubles the peak memory for long T-Rex episodes. Even with Flash Attention, CoTracker's feature/correlation tensors grow linearly with the frame count. Each new window is initialized from the previous trajectory, and overlapping predictions are blended to avoid a discontinuity at the boundary. """ video_np = np.asarray(video_hwc, dtype=np.uint8) total_frames = int(video_np.shape[0]) num_queries = int(query_xy.shape[0]) window_frames = int(os.environ.get("TREX_COTRACKER_WINDOW_FRAMES", "768")) overlap_frames = int(os.environ.get("TREX_COTRACKER_WINDOW_OVERLAP", "64")) if window_frames < 2: raise ValueError("TREX_COTRACKER_WINDOW_FRAMES must be at least 2") if overlap_frames < 1 or overlap_frames >= window_frames: raise ValueError( "TREX_COTRACKER_WINDOW_OVERLAP must be in [1, WINDOW_FRAMES)" ) if total_frames <= window_frames: return _run_cotracker_window(model, video_np, query_xy, device) step = window_frames - overlap_frames num_windows = 1 + (total_frames - window_frames + step - 1) // step print( f" CoTracker windowing: {total_frames} frames -> {num_windows} " f"window(s), max={window_frames}, overlap={overlap_frames}" ) tracks = np.empty((total_frames, num_queries, 2), dtype=np.float32) visibility = np.empty((total_frames, num_queries), dtype=np.float32) start = 0 filled_end = 0 while start < total_frames: end = min(start + window_frames, total_frames) seed_xy = np.asarray(query_xy if start == 0 else tracks[start], dtype=np.float32) window_tracks, window_visibility = _run_cotracker_window( model, video_np[start:end], seed_xy, device, ) overlap_end = min(filled_end, end) existing_frames = max(0, overlap_end - start) if existing_frames > 0: alpha = np.linspace( 0.0, 1.0, existing_frames, dtype=np.float32, ) tracks[start:overlap_end] = ( tracks[start:overlap_end] * (1.0 - alpha[:, None, None]) + window_tracks[:existing_frames] * alpha[:, None, None] ) use_new_visibility = alpha >= 0.5 visibility[start:overlap_end] = np.where( use_new_visibility[:, None], window_visibility[:existing_frames], visibility[start:overlap_end], ) tracks[overlap_end:end] = window_tracks[existing_frames:] visibility[overlap_end:end] = window_visibility[existing_frames:] filled_end = max(filled_end, end) if end >= total_frames: break start = end - overlap_frames return tracks, visibility def _tracks_episode( *, view_images: dict[str, np.ndarray], out_hw: tuple[int, int], cotracker_model, cotracker_device: object, sam2_predictor=None, sam2_seed: int | None = None, ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict[str, np.ndarray | None], dict[str, str]]: from trex_track.sam2_cotracker_hands import head_hands_50, wrist_hand_50 from trex_track.trex_projection import make_image_grid view_tracks_px: dict[str, np.ndarray] = {} view_vis: dict[str, np.ndarray] = {} masks: dict[str, np.ndarray | None] = {} tags: dict[str, str] = {} # Head: left hand/arm 50, then right hand/arm 50. head_q, left_m, right_m, head_tag = head_hands_50( sam2_predictor, view_images["head_left"][0], n_points=NUM_HAND_POINTS, seed=sam2_seed, ) if int(head_q.shape[0]) != NUM_HEAD_POINTS: raise ValueError(f"head queries expect {NUM_HEAD_POINTS}, got {head_q.shape[0]}") head_trk, head_vis = _run_cotracker( cotracker_model, view_images["head_left"], head_q, cotracker_device ) view_tracks_px["head_left"] = head_trk view_vis["head_left"] = head_vis masks["head_left_hand"] = left_m masks["head_right_hand"] = right_m tags["head_left"] = head_tag # Wrists: fixed 5×5 background grid + 50 SAM2 hand points. for view in ("left_wrist", "right_wrist"): hand_q, mask, tag = wrist_hand_50( sam2_predictor, view_images[view][0], view, n_points=NUM_WRIST_HAND, seed=sam2_seed, ) if int(hand_q.shape[0]) != NUM_WRIST_HAND: raise ValueError(f"{view} hand queries expect {NUM_WRIST_HAND}, got {hand_q.shape[0]}") grid_q = make_image_grid(out_hw[0], out_hw[1], grid_size=5).astype(np.float32) q = np.concatenate([grid_q, hand_q], axis=0) if int(q.shape[0]) != NUM_WRIST_POINTS: raise ValueError(f"{view} queries expect {NUM_WRIST_POINTS}, got {q.shape[0]}") trk, vis = _run_cotracker(cotracker_model, view_images[view], q, cotracker_device) view_tracks_px[view] = trk view_vis[view] = vis masks[view] = mask tags[view] = f"grid25+{tag}" view_tracks: dict[str, np.ndarray] = {} for view in VIEW_ORDER: view_tracks[view], view_vis[view] = normalize_track_result( view_tracks_px[view], view_vis[view], out_hw[1], out_hw[0], ) return view_tracks, view_vis, masks, tags def _save_seed_overlay( out_dir: Path, episode_index: int, view_images: dict[str, np.ndarray], masks: dict[str, np.ndarray | None], view_tracks: dict[str, np.ndarray], out_hw: tuple[int, int], ) -> None: """Save frame-0 overlays with mask + query points for debugging.""" import cv2 out_dir.mkdir(parents=True, exist_ok=True) h, w = out_hw def _draw(view: str, mask: np.ndarray | None, tracks_norm: np.ndarray, color_bgr): rgb = view_images[view][0] bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) if mask is not None and np.asarray(mask).any(): m = np.asarray(mask, dtype=bool) bgr[m] = (bgr[m] * 0.35 + np.asarray(color_bgr, dtype=np.float32) * 0.65).astype(np.uint8) pts = tracks_norm[0] * np.array([w, h], dtype=np.float32) for p in pts.astype(int): cv2.circle(bgr, (int(p[0]), int(p[1])), 3, (0, 255, 255), -1, cv2.LINE_AA) return bgr # Head: both masks head = cv2.cvtColor(view_images["head_left"][0], cv2.COLOR_RGB2BGR) lm, rm = masks.get("head_left_hand"), masks.get("head_right_hand") if lm is not None and np.asarray(lm).any(): m = np.asarray(lm, dtype=bool) head[m] = (head[m] * 0.35 + np.array([0, 255, 0], dtype=np.float32) * 0.65).astype(np.uint8) if rm is not None and np.asarray(rm).any(): m = np.asarray(rm, dtype=bool) head[m] = (head[m] * 0.35 + np.array([0, 165, 255], dtype=np.float32) * 0.65).astype(np.uint8) pts = view_tracks["head_left"][0] * np.array([w, h], dtype=np.float32) for i, p in enumerate(pts.astype(int)): col = (0, 255, 255) if i < NUM_HEAD_LEFT else (255, 255, 0) cv2.circle(head, (int(p[0]), int(p[1])), 3, col, -1, cv2.LINE_AA) cv2.imwrite(str(out_dir / f"episode_{episode_index:06d}_head_left_seeds.png"), head) for view in ("left_wrist", "right_wrist"): rgb = view_images[view][0] bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) mask = masks.get(view) if mask is not None and np.asarray(mask).any(): m = np.asarray(mask, dtype=bool) bgr[m] = (bgr[m] * 0.40 + np.array([0, 0, 220], dtype=np.float32) * 0.60).astype(np.uint8) pts = view_tracks[view][0] * np.array([w, h], dtype=np.float32) for i, p in enumerate(pts.astype(int)): # yellow = fixed background grid; magenta = force-aligned hand col = (0, 255, 255) if i < NUM_WRIST_GRID else (255, 0, 255) cv2.drawMarker( bgr, (int(p[0]), int(p[1])), col, markerType=cv2.MARKER_STAR, markerSize=8, thickness=1, line_type=cv2.LINE_AA, ) cv2.imwrite(str(out_dir / f"episode_{episode_index:06d}_{view}_seeds.png"), bgr) def _save_prompt_overlays( out_dir: Path, episode_index: int, view_images: dict[str, np.ndarray], masks: dict[str, np.ndarray | None], ) -> None: """Save SAM2 prompts (pos/neg/box) + mask overlays; reuse track masks (no 2nd SAM2).""" import cv2 from trex_track.sam2_prompt_hands import ( build_head_prompts, build_wrist_prompts, draw_prompt_overlay, ) out_dir.mkdir(parents=True, exist_ok=True) panels: list[np.ndarray] = [] meta_dump: dict[str, object] = {} def _one( key: str, rgb: np.ndarray, mask: np.ndarray | None, coords: np.ndarray, labels: np.ndarray, box: np.ndarray, mask_bgr: tuple[int, int, int], ) -> np.ndarray: vis = draw_prompt_overlay( rgb, mask=mask, coords=coords, labels=labels, box=box, mask_bgr=mask_bgr, ) # Legend n_pos = int((labels == 1).sum()) n_neg = int((labels == 0).sum()) area = int(np.asarray(mask).sum()) if mask is not None else 0 cv2.putText( vis, f"{key} +pos={n_pos} -neg={n_neg} mask_px={area}", (6, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.40, (255, 255, 255), 1, cv2.LINE_AA, ) path = out_dir / f"episode_{episode_index:06d}_{key}_prompt.png" cv2.imwrite(str(path), vis) print(f"Wrote {path}") meta_dump[f"{key}_coords"] = np.asarray(coords, dtype=np.float32) meta_dump[f"{key}_labels"] = np.asarray(labels, dtype=np.int32) meta_dump[f"{key}_box"] = np.asarray(box, dtype=np.float32) if mask is not None: meta_dump[f"{key}_mask"] = np.asarray(mask, dtype=bool) return vis # Wrist views for view, col in (("left_wrist", (0, 0, 220)), ("right_wrist", (0, 0, 220))): rgb = view_images[view][0] h, w = rgb.shape[:2] coords, labels, box = build_wrist_prompts(h, w, view) panels.append(_one(view, rgb, masks.get(view), coords, labels, box, col)) # Head: left / right hands (same RGB, two prompt sets) rgb = view_images["head_left"][0] h, w = rgb.shape[:2] for side, col, mkey in ( ("left", (0, 255, 0), "head_left_hand"), ("right", (0, 165, 255), "head_right_hand"), ): coords, labels, box = build_head_prompts(h, w, side) panels.append(_one(f"head_{side}", rgb, masks.get(mkey), coords, labels, box, col)) # Combined 2x2 panel for quick inspection if len(panels) == 4: top = np.concatenate(panels[2:4], axis=1) # head_left | head_right bot = np.concatenate(panels[0:2], axis=1) # left_wrist | right_wrist # Resize to same width if needed if top.shape[1] != bot.shape[1]: tw = max(top.shape[1], bot.shape[1]) top = cv2.resize(top, (tw, top.shape[0])) bot = cv2.resize(bot, (tw, bot.shape[0])) grid = np.concatenate([top, bot], axis=0) grid_path = out_dir / f"episode_{episode_index:06d}_all_prompts.png" cv2.imwrite(str(grid_path), grid) print(f"Wrote {grid_path}") meta_path = out_dir / f"episode_{episode_index:06d}_prompts.npz" np.savez_compressed(meta_path, **meta_dump) print(f"Wrote {meta_path}") def _atomic_savez(path: Path, **arrays: object) -> None: """Write an NPZ in the destination directory, then atomically replace.""" path.parent.mkdir(parents=True, exist_ok=True) fd, tmp_name = tempfile.mkstemp( prefix=f".{path.name}.", suffix=".tmp", dir=path.parent, ) try: with os.fdopen(fd, "wb") as file: np.savez_compressed(file, **arrays) file.flush() os.fsync(file.fileno()) os.replace(tmp_name, path) except BaseException: try: os.unlink(tmp_name) except FileNotFoundError: pass raise def process_episode( *, dataset_root: Path, episode_index: int, output_path: Path, calib: dict | None, out_hw: tuple[int, int], cotracker_model, cotracker_device: object, save_viz: bool, viz_out_dir: Path, viz_fps: int, viz_trail: int, sam2_predictor=None, sam2_seed: int | None = None, save_sam2_masks_flag: bool = True, sam2_masks_dir: Path | None = None, ) -> Path: from trex_track.sam2_wrist_hand import save_sam2_masks from trex_track.trex_viz_tracks import render_three_view_combined_video del calib # Reserved for provenance/backward-compatible callers. states, task = load_episode_states(dataset_root, episode_index) view_images = load_episode_videos(dataset_root, episode_index, out_hw=out_hw) frame_counts = {"parquet": int(states.shape[0])} frame_counts.update({view: int(view_images[view].shape[0]) for view in VIEW_ORDER}) if len(set(frame_counts.values())) != 1: raise ValueError( f"episode {episode_index}: frame-count mismatch; refusing to truncate: " f"{frame_counts}" ) t_len = int(states.shape[0]) view_tracks, view_vis, masks, tags = _tracks_episode( view_images=view_images, out_hw=out_hw, cotracker_model=cotracker_model, cotracker_device=cotracker_device, sam2_predictor=sam2_predictor, sam2_seed=sam2_seed, ) tracks_combined = np.concatenate([view_tracks[v] for v in VIEW_ORDER], axis=1) vis_combined = np.concatenate([view_vis[v] for v in VIEW_ORDER], axis=1) if tracks_combined.shape != (t_len, NUM_COMBINED_POINTS, 2): raise ValueError( f"episode {episode_index}: combined tracks have {tracks_combined.shape}, " f"expected {(t_len, NUM_COMBINED_POINTS, 2)}" ) if vis_combined.shape != (t_len, NUM_COMBINED_POINTS): raise ValueError( f"episode {episode_index}: combined visibility has {vis_combined.shape}" ) out_npz = output_path / f"episode_{episode_index:06d}.npz" identities = identity_metadata() _atomic_savez( out_npz, tracks=tracks_combined, vis=vis_combined, tracks_head_left=view_tracks["head_left"], tracks_left_wrist=view_tracks["left_wrist"], tracks_right_wrist=view_tracks["right_wrist"], vis_head_left=view_vis["head_left"], vis_left_wrist=view_vis["left_wrist"], vis_right_wrist=view_vis["right_wrist"], language=np.array(task), episode_index=np.array(episode_index, dtype=np.int32), num_steps=np.array(t_len, dtype=np.int32), point_slices=np.array(POINT_SLICES, dtype=np.int32), point_view_ids=np.asarray(identities["view_ids"], dtype=np.int8), point_hand_ids=np.asarray(identities["hand_ids"], dtype=np.int8), point_role_ids=np.asarray(identities["role_ids"], dtype=np.int8), point_local_ids=np.asarray(identities["local_ids"], dtype=np.int16), point_global_ids=np.asarray(identities["global_ids"], dtype=np.int16), point_names=np.asarray(identities["point_names"]), points_per_hand=np.array(NUM_HAND_POINTS, dtype=np.int32), wrist_grid_points=np.array(NUM_WRIST_GRID, dtype=np.int32), wrist_hand_points=np.array(NUM_WRIST_HAND, dtype=np.int32), head_query_source=np.array(tags.get("head_left", "sam2")), left_wrist_query_source=np.array(tags.get("left_wrist", "sam2")), right_wrist_query_source=np.array(tags.get("right_wrist", "sam2")), track_source=np.array("sam2_once_prompt_cotracker"), tracks_coord_space=np.array("normalized_div_wh"), track_layout_version=np.array(TRACK_LAYOUT_VERSION), proj_image_hw=np.array(out_hw, dtype=np.int32), ) masks_dir = sam2_masks_dir if sam2_masks_dir is not None else output_path / "sam2_masks" if save_sam2_masks_flag: # Save binary masks for wrist + head hands save_map = { "left_wrist": masks.get("left_wrist"), "right_wrist": masks.get("right_wrist"), "head_left_hand": masks.get("head_left_hand"), "head_right_hand": masks.get("head_right_hand"), } for p in save_sam2_masks(masks_dir, episode_index, save_map): print(f"Wrote {p}") _save_seed_overlay(masks_dir, episode_index, view_images, masks, view_tracks, out_hw) _save_prompt_overlays(masks_dir, episode_index, view_images, masks) if save_viz: render_three_view_combined_video( view_images=view_images, view_tracks=view_tracks, view_vis=view_vis, out_path=viz_out_dir / f"episode_{episode_index:06d}.mp4", fps=viz_fps, draw_trail=viz_trail, ) print(f"Wrote {out_npz}") if save_viz: print(f"Wrote {viz_out_dir / f'episode_{episode_index:06d}.mp4'}") print(f" tags: {tags}") return out_npz def create_tracking_runtime( *, calib_path: str | Path = DEFAULT_CALIB, openpi_root: str | Path = DEFAULT_OPENPI_ROOT, cotracker_checkpoint: str | Path | None = None, cotracker_device: str = "", sam2_model: str = DEFAULT_SAM2_MODEL, sam2_device: str = "", sam2_libs: str | Path = DEFAULT_SAM2_LIBS, image_height: int = 0, image_width: int = 0, ) -> TrackingRuntime: """Load CoTracker and SAM2 once; safe to call from the batch builder.""" import torch from trex_track.sam2_wrist_hand import load_sam2_predictor from trex_track.trex_projection import load_camera_calib calib = load_camera_calib(calib_path) if bool(image_height > 0) != bool(image_width > 0): raise ValueError("image-height and image-width must be set together") if image_height > 0: out_hw = (int(image_height), int(image_width)) else: out_hw = tuple(int(x) for x in calib.get("video_hw", [180, 320])) _ensure_openpi_on_path(Path(openpi_root)) from utils.cotracker_wrist_grid import ( # type: ignore default_cotracker_checkpoint, load_cotracker_predictor, ) _enable_cotracker_sdpa_attention(openpi_root) checkpoint_text = str(cotracker_checkpoint or "").strip() checkpoint = ( Path(checkpoint_text).expanduser().resolve() if checkpoint_text else default_cotracker_checkpoint() ) device_text = cotracker_device.strip() or ( "cuda:0" if torch.cuda.is_available() else "cpu" ) device = torch.device(device_text) print(f" CoTracker: {checkpoint} on {device}") if device.type == "cuda": print(" CoTracker attention: BF16 PyTorch SDPA (Flash-compatible)") cotracker_model = load_cotracker_predictor(checkpoint, device) sam_device = sam2_device.strip() or device_text print(f" SAM2: {sam2_model} on {sam_device}") sam2_predictor = load_sam2_predictor( model_id=str(sam2_model), device=sam_device, sam2_libs=sam2_libs, ) return TrackingRuntime( calib=calib, out_hw=out_hw, cotracker_model=cotracker_model, cotracker_device=device, sam2_predictor=sam2_predictor, ) def main() -> None: parser = argparse.ArgumentParser( description="Extract canonical T-Rex SAM2+CoTracker tracks (250 points)" ) parser.add_argument( "--dataset-root", type=str, default=str(_DREAMZERO_ROOT / "data" / "trex_small"), ) parser.add_argument("--episode-index", type=int, default=0) parser.add_argument( "--output-path", type=str, default=str(_DREAMZERO_ROOT / "data" / "trex_small_tracks"), ) parser.add_argument("--calib-path", type=str, default=str(DEFAULT_CALIB)) parser.add_argument("--openpi-root", type=str, default=str(DEFAULT_OPENPI_ROOT)) parser.add_argument("--cotracker-checkpoint", type=str, default="") parser.add_argument("--cotracker-device", type=str, default="") parser.add_argument("--image-height", type=int, default=0) parser.add_argument("--image-width", type=int, default=0) parser.add_argument("--save-viz", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--viz-out-dir", type=str, default="") parser.add_argument("--viz-fps", type=int, default=10) parser.add_argument("--viz-trail", type=int, default=15) parser.add_argument( "--sam2-model", type=str, default=DEFAULT_SAM2_MODEL, ) parser.add_argument("--sam2-device", type=str, default="") parser.add_argument("--sam2-seed", type=int, default=0) parser.add_argument("--save-sam2-masks", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--sam2-masks-dir", type=str, default="") parser.add_argument( "--sam2-libs", type=str, default=DEFAULT_SAM2_LIBS, ) args = parser.parse_args() dataset_root = Path(args.dataset_root).expanduser().resolve() output_path = Path(args.output_path).expanduser().resolve() runtime = create_tracking_runtime( calib_path=args.calib_path, openpi_root=args.openpi_root, cotracker_checkpoint=args.cotracker_checkpoint, cotracker_device=args.cotracker_device, sam2_model=args.sam2_model, sam2_device=args.sam2_device, sam2_libs=args.sam2_libs, image_height=args.image_height, image_width=args.image_width, ) viz_out_dir = ( Path(args.viz_out_dir).expanduser().resolve() if args.viz_out_dir.strip() else output_path / "viz_tracks" ) sam2_masks_dir = ( Path(args.sam2_masks_dir).expanduser().resolve() if args.sam2_masks_dir.strip() else output_path / "sam2_masks" ) print("T-Rex SAM2+CoTracker extraction") print(f" Dataset: {dataset_root}") print(f" Episode: {args.episode_index}") print(f" Output: {output_path}") print(f" Image: {runtime.out_hw[1]}x{runtime.out_hw[0]}") print( f" Points: head={NUM_HEAD_POINTS} (50+50), " f"wrist={NUM_WRIST_POINTS}x2 (grid{NUM_WRIST_GRID}+hand{NUM_WRIST_HAND}), " f"total={NUM_COMBINED_POINTS}" ) process_episode( dataset_root=dataset_root, episode_index=int(args.episode_index), output_path=output_path, calib=runtime.calib, out_hw=runtime.out_hw, cotracker_model=runtime.cotracker_model, cotracker_device=runtime.cotracker_device, save_viz=bool(args.save_viz), viz_out_dir=viz_out_dir, viz_fps=int(args.viz_fps), viz_trail=int(args.viz_trail), sam2_predictor=runtime.sam2_predictor, sam2_seed=int(args.sam2_seed), save_sam2_masks_flag=bool(args.save_sam2_masks), sam2_masks_dir=sam2_masks_dir, ) if __name__ == "__main__": main()