scripts/extract_track.py

"""
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()