scripts/measure_push_tracked.py

#!/usr/bin/env python
"""Measure a real push using TAPNext++ point tracks (not a mask centroid).

Supersedes `scripts/measure_push_demo.py`'s centroid method, which failed on
this episode for a concrete, measured reason: the arm occludes part of the
drawer while pushing it, so the visible mask's centroid drifts with the
*occlusion pattern* rather than with the drawer. That produced a ~40-100 px
systematic reprojection residual (~4-11 cm at this scene's ~9 px/cm) and a
displacement profile that rose and then returned to zero -- physically
impossible for a drawer being closed. Occlusion-robust point tracks fix it.

Method:
  1. Seed query points inside the drawer's mask on the depth frame (frame 0),
     where each point's metric 3D position is directly available.
  2. Track them with TAPNext++ across the whole window up to and through the
     push. Points the arm covers are reported not-visible and simply drop out
     of that frame's fit rather than corrupting it.
  3. Per frame, solve one rigid horizontal translation (dx, dy) shared by all
     visible points, by least squares over their reprojections. A drawer runs
     level, so dz = 0 is a physical constraint, not an approximation; using
     many points instead of one centroid makes the fit over-determined and lets
     the residual actually report whether the rigid-slide model holds.

Usage:
    PYTHONPATH=src python scripts/measure_push_tracked.py \\
        --episode <uuid> --mask-outputs outputs/<uuid>/objects_shelf2/0_11 \\
        --traj-end 124 --out outputs/push_calibration_tracked.json
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import sys
from pathlib import Path

import h5py
import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from fpgm.config import PipelineConfig  # noqa: E402
from fpgm.data.droid_raw import cam2world_vector_to_world2cam  # noqa: E402
from fpgm.geometry.camera import Camera  # noqa: E402
from fpgm.pipeline.frames import ClipFrameSource  # noqa: E402
from fpgm.robot.urdf import RobotModel  # noqa: E402
from fpgm.tracking.sampling import sample_points_in_mask  # noqa: E402
from fpgm.tracking.tapnext import TapNextPointTracker  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

_spec = importlib.util.spec_from_file_location(
    "run_object_pipeline", REPO_ROOT / "scripts" / "run_object_pipeline.py"
)
rop = importlib.util.module_from_spec(_spec)
sys.modules[_spec.name] = rop
_spec.loader.exec_module(rop)

logger = get_logger("measure_push_tracked")

_TRAJ_FPS = 15.0


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--episode", required=True)
    p.add_argument("--camera", default="ext1", choices=["ext1", "ext2"])
    p.add_argument("--mask-outputs", type=Path, required=True)
    p.add_argument("--clip", default="0:11")
    p.add_argument("--traj-end", type=int, default=124)
    p.add_argument("--n-points", type=int, default=96)
    p.add_argument("--device", default="cuda")
    p.add_argument("--out", type=Path, required=True)
    return p.parse_args()


def _lift_points_to_3d(mask, uv, adapter, camera_annot, frame_shape):
    """Metric 3D position for each query pixel, from the depth frame."""
    depth_m = adapter.initial_depth.astype(np.float64) / 1000.0
    dh, dw = depth_m.shape
    fh, fw = frame_shape
    # Query points are in video pixels; depth is at annotation resolution.
    su, sv = dw / fw, dh / fh
    du = np.clip(np.round(uv[:, 0] * su).astype(int), 0, dw - 1)
    dv = np.clip(np.round(uv[:, 1] * sv).astype(int), 0, dh - 1)
    z = depth_m[dv, du]
    ok = z > 0
    cam = camera_annot
    uv_annot = np.stack([uv[:, 0] * su, uv[:, 1] * sv], axis=1)
    pts = cam.unproject(uv_annot[ok], z[ok])
    return pts, ok


def main() -> int:
    args = parse_args()
    setup_logging()

    meta = json.loads((rop.stage_dir(args.mask_outputs, "mask") / "meta.json").read_text())
    with np.load(rop.stage_dir(args.mask_outputs, "mask") / "mask_frame0.npz") as npz:
        mask = npz["mask"].astype(bool)
    video_w, video_h = meta["video_resolution"]

    ctx = rop._load_clip_context(args.episode, args.camera, args.clip)
    depth = rop._load_initial_depth(args.episode, args.clip, meta["camera_serial"])
    adapter = rop._DepthRgbAdapter(depth, ctx.clip.initial_rgb)

    uv_seed = sample_points_in_mask(mask, args.n_points, method="farthest", boundary_erosion_px=6)
    X_seed, ok = _lift_points_to_3d(mask, uv_seed, adapter, ctx.camera_annot, (video_h, video_w))
    uv_seed = uv_seed[ok]
    logger.info("seeded %d points with valid depth (of %d sampled)", X_seed.shape[0], ok.size)
    if X_seed.shape[0] < 8:
        raise SystemExit("too few query points survived the depth check")

    with h5py.File(rop._load_trajectory_h5(args.episode), "r") as f:
        J = np.asarray(f["observation/robot_state/joint_positions"])
        G = np.asarray(f["observation/robot_state/gripper_position"])
        extrinsics_6d = np.asarray(f[f"observation/camera_extrinsics/{meta['camera_serial']}_left"])

    n_frames = args.traj_end + 1
    cfg = PipelineConfig()
    tracker = TapNextPointTracker(
        cfg.tracking, checkpoint_path=str(cfg.paths.checkpoints_dir / cfg.tracking.checkpoint),
        device=args.device,
    )
    with ClipFrameSource(str(ctx.mp4_path), 0, n_frames, stride=1) as frames:
        rgb = (f[:, :, ::-1] for f in (frames.read(i) for i in range(frames.n_frames)))
        track = tracker.track(rgb, uv_seed.astype(np.float32), query_frame_idx=0)

    uv_t = np.asarray(track.uv)            # (T, Q, 2)
    vis_t = np.asarray(track.visible)      # (T, Q)
    n_t = uv_t.shape[0]
    logger.info("tracked %d frames x %d points", n_t, uv_t.shape[1])

    K = ctx.camera_annot.rescaled(video_w, video_h).K
    grid = np.arange(-0.30, 0.30 + 1e-9, 0.0025)
    dx, dy = np.meshgrid(grid, grid, indexing="ij")
    offsets = np.stack([dx.ravel(), dy.ravel(), np.zeros(dx.size)], axis=1)

    # --- keep only points that actually ride the drawer ----------------------
    # The SAM mask of an *open* drawer necessarily also covers its contents and
    # slivers of the static cabinet behind it. Those points do not translate
    # with the drawer, so including them biases the shared-translation fit and
    # is exactly what drives the reprojection residual up as the drawer moves.
    # Each point is given its own 2-DOF solve on the last frame (2 unknowns, 2
    # observations -- exactly determined); the drawer's points agree on one
    # displacement while static ones cluster at zero, so the majority cluster
    # around the median identifies the drawer.
    last = n_t - 1
    cam_last = Camera(K, cam2world_vector_to_world2cam(extrinsics_6d[last]))
    per_point = np.full((X_seed.shape[0], 2), np.nan)
    vis_last = vis_t[last] & np.isfinite(uv_t[last]).all(axis=1)
    for q in np.nonzero(vis_last)[0]:
        cand = X_seed[q][None, :] + offsets
        uvp, dep = cam_last.project(cand)
        err = np.linalg.norm(uvp - uv_t[last][q][None, :], axis=1)
        err[dep <= 0] = 1e6
        per_point[q] = offsets[int(np.argmin(err)), :2]

    finite = np.isfinite(per_point[:, 0])
    med = np.median(per_point[finite], axis=0)
    dev = np.linalg.norm(per_point - med[None, :], axis=1)
    mad = np.median(dev[finite])
    inlier = finite & (dev <= max(0.02, 2.0 * mad))
    logger.info(
        "drawer-point selection: %d/%d inliers (median displacement %s m, MAD %.3f m)",
        int(inlier.sum()), int(finite.sum()), med.round(3), float(mad),
    )
    if inlier.sum() < 8:
        raise SystemExit("too few consistent drawer points -- rigid-slide model does not hold")
    X_seed = X_seed[inlier]
    uv_t = uv_t[:, inlier, :]
    vis_t = vis_t[:, inlier]

    disp_xy = np.full((n_t, 2), np.nan)
    resid = np.full(n_t, np.nan)
    n_vis = np.zeros(n_t, dtype=int)
    for t in range(n_t):
        vis = vis_t[t] & np.isfinite(uv_t[t]).all(axis=1)
        n_vis[t] = int(vis.sum())
        if n_vis[t] < 8:
            continue
        camera = Camera(K, cam2world_vector_to_world2cam(extrinsics_6d[t]))
        # Shared rigid horizontal translation: evaluate all candidate offsets
        # against every visible point at once and keep the least-squares best.
        best_err, best_off = np.inf, None
        Xv = X_seed[vis]
        obs = uv_t[t][vis]
        for chunk in np.array_split(offsets, 8):
            cand = (Xv[None, :, :] + chunk[:, None, :]).reshape(-1, 3)
            uvp, dep = camera.project(cand)
            uvp = uvp.reshape(chunk.shape[0], Xv.shape[0], 2)
            dep = dep.reshape(chunk.shape[0], Xv.shape[0])
            err = np.linalg.norm(uvp - obs[None, :, :], axis=2)
            err[dep <= 0] = 1e6
            rms = np.sqrt((err ** 2).mean(axis=1))
            i = int(np.argmin(rms))
            if rms[i] < best_err:
                best_err, best_off = float(rms[i]), chunk[i, :2].copy()
        disp_xy[t] = best_off
        resid[t] = best_err

    valid = np.isfinite(disp_xy[:, 0])
    ref = np.nanmedian(disp_xy[:20][np.isfinite(disp_xy[:20, 0])], axis=0)
    disp_xy = disp_xy - ref[None, :]

    robot = RobotModel(str(rop._DEFAULT_URDF), load_meshes=False)
    hand = np.array([
        (robot.link_poses(J[t], float(G[t]))["left_inner_finger"][:3, 3]
         + robot.link_poses(J[t], float(G[t]))["right_inner_finger"][:3, 3]) / 2.0
        for t in range(n_t)
    ])

    result = {
        "episode": args.episode,
        # The inlier seeds' 3D positions ARE the drawer, established by the data
        # rather than assumed: they are exactly the points that translated
        # rigidly together throughout the push. Downstream geometry should be
        # fitted to these, not to the raw SAM mask's point cloud (which also
        # covers the drawer's contents and the static cabinet behind it).
        "drawer_inlier_points_world": X_seed.tolist(),
        "n_query_points": int(X_seed.shape[0]),
        "n_frames": int(n_t),
        "displacement_xy_m": np.where(np.isfinite(disp_xy), disp_xy, None).tolist(),
        "residual_px": np.where(np.isfinite(resid), resid, None).tolist(),
        "n_visible_per_frame": n_vis.tolist(),
        "hand_xyz_m": hand.tolist(),
        "traj_fps": _TRAJ_FPS,
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(result, indent=2))

    mag = np.linalg.norm(disp_xy, axis=1)
    print(f"{'row':>4} {'dx':>8} {'dy':>8} {'|d|':>8} {'resid_px':>9} {'n_vis':>6}")
    for t in range(80, n_t):
        if not valid[t]:
            continue
        print(f"{t:4d} {disp_xy[t,0]:8.3f} {disp_xy[t,1]:8.3f} {mag[t]:8.3f} {resid[t]:9.1f} {n_vis[t]:6d}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())