scripts/trex_track/trex_projection.py

"""Project T-Rex hand keypoints with camera intrinsics / extrinsics."""

from __future__ import annotations

import json
from pathlib import Path

import cv2
import numpy as np
from scipy.spatial.transform import Rotation

EULER_CONVENTION = "xyz"
TREX_NATIVE_HW = (360, 640)  # (H, W) — robot stream resolution before v2 downscale


def pose6_to_matrix(pose6: np.ndarray, *, euler: str = EULER_CONVENTION) -> np.ndarray:
    x = np.asarray(pose6, dtype=np.float64).reshape(-1)
    rot = Rotation.from_euler(euler, x[3:6]).as_matrix()
    T = np.eye(4, dtype=np.float64)
    T[:3, :3] = rot
    T[:3, 3] = x[:3]
    return T


def matrix_to_pose6(T: np.ndarray, *, euler: str = EULER_CONVENTION) -> np.ndarray:
    T = np.asarray(T, dtype=np.float64).reshape(4, 4)
    xyz = T[:3, 3]
    rpy = Rotation.from_matrix(T[:3, :3]).as_euler(euler)
    return np.concatenate([xyz, rpy], axis=0)


def scale_K_to_image_hw(
    K: np.ndarray,
    native_hw: tuple[int, int],
    image_hw: tuple[int, int],
) -> np.ndarray:
    href, wref = native_hw
    hnew, wnew = image_hw
    K = np.asarray(K, dtype=np.float64).reshape(3, 3).copy()
    K[0, 0] *= wnew / float(wref)
    K[0, 2] *= wnew / float(wref)
    K[1, 1] *= hnew / float(href)
    K[1, 2] *= hnew / float(href)
    return K


def scale_pixels_between_hw(
    pixels: np.ndarray,
    src_hw: tuple[int, int],
    dst_hw: tuple[int, int],
) -> np.ndarray:
    out = np.asarray(pixels, dtype=np.float64).copy()
    sh, sw = src_hw
    dh, dw = dst_hw
    out[..., 0] *= dw / float(sw)
    out[..., 1] *= dh / float(sh)
    return out


def make_image_grid(img_h: int, img_w: int, grid_size: int = 5) -> np.ndarray:
    xs = np.linspace(0.0, float(img_w - 1), grid_size, dtype=np.float64)
    ys = np.linspace(0.0, float(img_h - 1), grid_size, dtype=np.float64)
    return np.stack(np.meshgrid(xs, ys, indexing="xy"), axis=-1).reshape(-1, 2)


# Canonical 22-joint hand layout in a unit box (x,y in [0,1], palm down, fingers up).
_HAND_TEMPLATE_UNIT = np.array(
    [
        [0.18, 0.82],
        [0.22, 0.72],
        [0.26, 0.62],
        [0.30, 0.52],
        [0.34, 0.42],
        [0.42, 0.88],
        [0.46, 0.74],
        [0.50, 0.60],
        [0.54, 0.46],
        [0.50, 0.90],
        [0.54, 0.76],
        [0.58, 0.62],
        [0.62, 0.48],
        [0.58, 0.88],
        [0.62, 0.74],
        [0.66, 0.60],
        [0.70, 0.46],
        [0.72, 0.86],
        [0.76, 0.72],
        [0.80, 0.58],
        [0.84, 0.44],
        [0.88, 0.34],
    ],
    dtype=np.float64,
)

# T-Rex wrist cameras: hand sits in the lower band; left/right are mirrored in x.
# Anchors tuned on trex_small episode_000000 frame-0 hand blob.
_WRIST_HAND_VIEW_ANCHORS = {
    "left_wrist": {"cx": 0.60, "cy": 0.66, "sx": 0.36, "sy": 0.44, "mirror_x": False},
    "right_wrist": {"cx": 0.50, "cy": 0.66, "sx": 0.38, "sy": 0.46, "mirror_x": True},
}


def _build_wrist_hand_template_frac(view: str) -> np.ndarray:
    if view not in _WRIST_HAND_VIEW_ANCHORS:
        raise KeyError(f"unknown wrist view {view!r}")
    cfg = _WRIST_HAND_VIEW_ANCHORS[view]
    pts = _HAND_TEMPLATE_UNIT.copy()
    if cfg["mirror_x"]:
        pts[:, 0] = 1.0 - pts[:, 0]
    pts[:, 0] = (pts[:, 0] - 0.5) * cfg["sx"] + cfg["cx"]
    pts[:, 1] = (pts[:, 1] - 0.5) * cfg["sy"] + cfg["cy"]
    pts[:, 0] = np.clip(pts[:, 0], 0.02, 0.98)
    pts[:, 1] = np.clip(pts[:, 1], 0.02, 0.98)
    return pts


WRIST_HAND_TEMPLATE_FRAC_BY_VIEW = {
    view: _build_wrist_hand_template_frac(view) for view in _WRIST_HAND_VIEW_ANCHORS
}
# Back-compat alias (left wrist).
WRIST_HAND_TEMPLATE_FRAC = WRIST_HAND_TEMPLATE_FRAC_BY_VIEW["left_wrist"]
NUM_WRIST_HAND_TEMPLATE = int(_HAND_TEMPLATE_UNIT.shape[0])

# FK seeds are only trusted when enough points land in the lower hand band.
WRIST_HAND_FK_MIN_VALID = 8
WRIST_HAND_FK_MIN_MEAN_Y_FRAC = 0.45


def wrist_hand_template_pixels(image_hw: tuple[int, int], view: str = "left_wrist") -> np.ndarray:
    h, w = image_hw
    frac = WRIST_HAND_TEMPLATE_FRAC_BY_VIEW[view]
    return (frac * np.array([float(w), float(h)], dtype=np.float64)).astype(np.float64)


def wrist_hand_prompt_box(image_hw: tuple[int, int], view: str = "left_wrist") -> np.ndarray:
    """XYXY box prompt for SAM2, derived from per-view hand anchors."""
    if view not in _WRIST_HAND_VIEW_ANCHORS:
        raise KeyError(f"unknown wrist view {view!r}")
    h, w = image_hw
    cfg = _WRIST_HAND_VIEW_ANCHORS[view]
    cx = float(cfg["cx"]) * float(w)
    cy = float(cfg["cy"]) * float(h)
    half_w = float(cfg["sx"]) * float(w) * 0.5
    half_h = float(cfg["sy"]) * float(h) * 0.5
    x0 = np.clip(cx - half_w, 0.0, float(w - 1))
    y0 = np.clip(cy - half_h, 0.0, float(h - 1))
    x1 = np.clip(cx + half_w, 0.0, float(w - 1))
    y1 = np.clip(cy + half_h, 0.0, float(h - 1))
    if x1 <= x0:
        x1 = min(float(w - 1), x0 + 1.0)
    if y1 <= y0:
        y1 = min(float(h - 1), y0 + 1.0)
    return np.array([x0, y0, x1, y1], dtype=np.float32)


def fk_hand_queries_trusted(
    queries: np.ndarray,
    valid: np.ndarray,
    image_hw: tuple[int, int],
) -> bool:
    """Reject wrist FK seeds that sit in the background instead of on the hand."""
    valid = np.asarray(valid, dtype=bool).reshape(-1)
    if int(valid.sum()) < WRIST_HAND_FK_MIN_VALID:
        return False
    h = max(float(image_hw[0]), 1.0)
    mean_y_frac = float(np.asarray(queries, dtype=np.float32)[valid, 1].mean() / h)
    return mean_y_frac >= WRIST_HAND_FK_MIN_MEAN_Y_FRAC


def queries_in_image(
    pixels: np.ndarray,
    vis: np.ndarray,
    image_hw: tuple[int, int],
) -> np.ndarray:
    h, w = image_hw
    pts = np.asarray(pixels, dtype=np.float64)
    mask = np.asarray(vis, dtype=bool).reshape(-1)
    mask &= np.isfinite(pts).all(axis=1)
    mask &= (pts[:, 0] >= 0.0) & (pts[:, 0] < float(w))
    mask &= (pts[:, 1] >= 0.0) & (pts[:, 1] < float(h))
    return mask


def align_hand_queries_to_template(
    queries: np.ndarray,
    valid: np.ndarray,
    image_hw: tuple[int, int],
    *,
    view: str = "left_wrist",
) -> np.ndarray:
    """Keep valid frame-0 projections; place invalid joints from a hand template."""
    out = wrist_hand_template_pixels(image_hw, view=view).astype(np.float32, copy=True)
    valid = np.asarray(valid, dtype=bool).reshape(-1)
    q = np.asarray(queries, dtype=np.float32)
    if valid.any():
        vpts = q[valid]
        vmin, vmax = vpts.min(axis=0), vpts.max(axis=0)
        tmpl = out.copy()
        tmin, tmax = tmpl.min(axis=0), tmpl.max(axis=0)
        scale = (vmax - vmin) / np.maximum(tmax - tmin, 1e-3)
        out = (tmpl - tmin) * scale + vmin
        out[valid] = q[valid]
    out[..., 0] = np.clip(out[..., 0], 0.0, float(image_hw[1] - 1))
    out[..., 1] = np.clip(out[..., 1], 0.0, float(image_hw[0] - 1))
    return out


def project_world_points(
    world_pts: np.ndarray,
    cam_to_world: np.ndarray,
    *,
    K: np.ndarray,
    dist: np.ndarray | None = None,
    native_hw: tuple[int, int] = TREX_NATIVE_HW,
    out_hw: tuple[int, int],
) -> tuple[np.ndarray, np.ndarray]:
    """Project (N,3) world points -> (N,2) pixels and visibility mask."""
    world = np.asarray(world_pts, dtype=np.float64).reshape(-1, 3)
    T_c2w = np.asarray(cam_to_world, dtype=np.float64).reshape(4, 4)
    w2c = np.linalg.inv(T_c2w)
    R = w2c[:3, :3]
    t = w2c[:3, 3].reshape(3, 1)

    K_native = np.asarray(K, dtype=np.float64).reshape(3, 3)
    dist_use = None
    if dist is not None:
        dist_use = np.asarray(dist, dtype=np.float64).reshape(-1)
        if dist_use.size >= 5 and np.allclose(dist_use[:5], 0.0):
            dist_use = None

    rvec, _ = cv2.Rodrigues(R)
    if dist_use is not None:
        pts2d, _ = cv2.projectPoints(
            world.reshape(-1, 1, 3),
            rvec,
            t,
            K_native,
            dist_use[:5],
        )
    else:
        pts2d, _ = cv2.projectPoints(world.reshape(-1, 1, 3), rvec, t, K_native, np.zeros(5))
    pixels_native = pts2d.reshape(-1, 2).astype(np.float64)

    cam_z = (R @ world.T + t).T[:, 2]
    vis = cam_z > 0.01

    pixels_out = scale_pixels_between_hw(pixels_native, native_hw, out_hw)
    href, wref = native_hw
    dh, dw = out_hw
    vis &= (
        (pixels_native[:, 0] >= 0.0)
        & (pixels_native[:, 0] < float(wref))
        & (pixels_native[:, 1] >= 0.0)
        & (pixels_native[:, 1] < float(href))
    )
    vis &= (
        (pixels_out[:, 0] >= 0.0)
        & (pixels_out[:, 0] < float(dw))
        & (pixels_out[:, 1] >= 0.0)
        & (pixels_out[:, 1] < float(dh))
    )
    return pixels_out.astype(np.float32), vis.astype(np.float32)


def wrist_cam_to_world(T_ee: np.ndarray, T_ee_to_cam: np.ndarray) -> np.ndarray:
    return np.asarray(T_ee, dtype=np.float64) @ np.asarray(T_ee_to_cam, dtype=np.float64)


def load_camera_calib(path: str | Path) -> dict:
    p = Path(path).expanduser().resolve()
    with p.open("r", encoding="utf-8") as f:
        return json.load(f)


def view_intrinsics(
    calib: dict,
    view: str,
    out_hw: tuple[int, int],
) -> tuple[np.ndarray, np.ndarray | None]:
    native_hw = tuple(calib.get("native_hw", TREX_NATIVE_HW))
    view_cfg = calib["views"][view]
    K = np.asarray(view_cfg["K_native"], dtype=np.float64).reshape(3, 3)
    dist = view_cfg.get("dist")
    if dist is not None:
        dist = np.asarray(dist, dtype=np.float64)
    K_out = scale_K_to_image_hw(K, native_hw, out_hw)
    return K_out, dist


def view_ee_to_cam(calib: dict, view: str) -> np.ndarray | None:
    view_cfg = calib["views"][view]
    if "T_ee_to_cam" in view_cfg:
        return np.asarray(view_cfg["T_ee_to_cam"], dtype=np.float64).reshape(4, 4)
    if "ee_to_cam_pose6" in view_cfg:
        return pose6_to_matrix(np.asarray(view_cfg["ee_to_cam_pose6"], dtype=np.float64))
    return None