| |
| """Push-Box point-flow overlay: GT CoTracker tracks vs 3-PoinTr predicted flow. |
| |
| GREEN = actual moving-points (CoTracker on the demo, the GT tracking points). |
| RED = 3-PoinTr predicted future trajectory, rolled out from frame-0's GT points. |
| |
| The push-box 3-PoinTr (push_box_2d_150pts_384) was trained on ANISOTROPICALLY |
| squashed 384x384 tracks, so we run inference in 384-space (rescale frame-0 GT |
| 1280x720 -> 384), then rescale the predicted flow back to native 1280x720 to draw |
| a crisp overlay on the original demo video. horizon L=249 (point_flow_head=498=249*2). |
| |
| Usage: |
| python gen_pred_flow_pushbox.py --ep ep000 [--ep ep005 ...] # or --all |
| """ |
| import argparse, os, sys |
| import cv2, imageio |
| import numpy as np |
| import torch |
|
|
| POINTR_REPO = "/home/hongyi/PWM/reward_model/3PoinTr_repo" |
| if POINTR_REPO not in sys.path: |
| sys.path.insert(0, POINTR_REPO) |
| from util.training_utils.arch import TrajectoryModel |
| from util.geometry_utils.geom_utils import ( |
| compute_start_sphere_params, sim_transform, sim_transform_inverse_points, |
| ) |
|
|
| CKPT = "/home/hongyi/PWM/reward_model/3pointr_models/push_box_2d_150pts_384_flow_backbone_final.pth" |
| TRACK_CACHE = "/data0/hongyi/reward_model_data/track_cache/demos_push_box" |
| PROMPT_SLUG = "cheez-it_box" |
| SQUASH = 384 |
| MAX_POINTS = 150 |
|
|
|
|
| def infer_horizon(state): |
| w = state.get("decoder_only.point_flow_head.weight") |
| return None if w is None else int(w.shape[0] // 2) |
|
|
|
|
| def load_model(ckpt_path, device): |
| ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) |
| cfg = ckpt.get("args", {}) or {} |
| state = ckpt.get("flow_backbone_state_dict") or ckpt.get("model_state_dict") or ckpt |
| horizon = infer_horizon(state) or int(cfg.get("horizon", 167)) |
| model = TrajectoryModel( |
| arch=cfg.get("arch", "flow_conditioned"), |
| dim=int(cfg.get("arch_dim", 256)), depth=int(cfg.get("arch_depth", 2)), |
| query_depth=int(cfg.get("arch_query_depth", 1)), |
| num_heads=int(cfg.get("arch_heads", 4)), |
| query_heads=int(cfg.get("arch_query_heads", 1)), |
| use_task_conditioning=False, num_tasks=1, horizon=horizon, |
| num_queries=int(cfg.get("num_queries", 10)), point_dim=2, |
| ).to(device) |
| missing, unexpected = model.load_state_dict(state, strict=False) |
| print(f"[predict] horizon={horizon} missing={len(missing)} unexpected={len(unexpected)}") |
| model.eval() |
| return model, horizon |
|
|
|
|
| @torch.no_grad() |
| def predict(model, horizon, x0_px, device): |
| """x0_px (N,2) in 384-space -> (L,N,2) predicted positions in 384-space.""" |
| x0 = torch.from_numpy(x0_px.astype(np.float32)).to(device).unsqueeze(0) |
| scales, center = compute_start_sphere_params(x0) |
| inv_scales = 1.0 / scales |
| x0n = sim_transform(x0, scales=inv_scales, rotations=None, translations=-center) |
| out = model({"x0": x0n, "num_steps": horizon, |
| "task_ids": torch.zeros(1, dtype=torch.long, device=device), |
| "scales": scales, "center": center}, mode="eval") |
| pred = out["pred_pos"] |
| L, N = pred.shape[1], pred.shape[2] |
| un = sim_transform_inverse_points(pred.reshape(1, L * N, 2), |
| scales=inv_scales, rotations=None, translations=-center) |
| return un.view(L, N, 2).cpu().numpy() |
|
|
|
|
| def render(ep, model, horizon, device, out_dir): |
| cache = f"{TRACK_CACHE}/tableHF_{ep}" |
| cap = cv2.VideoCapture(f"{cache}/input_short.mp4") |
| frames = [] |
| while True: |
| ok, f = cap.read() |
| if not ok: |
| break |
| frames.append(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) |
| cap.release() |
| Td = len(frames) |
| H, W = frames[0].shape[:2] |
| z = np.load(f"{cache}/{PROMPT_SLUG}/tracks_filtered.npz") |
| actual = z["tracks"].astype(np.float32) |
| vis = z["vis"] if "vis" in z.files else np.ones(actual.shape[:2], bool) |
| N = actual.shape[1] |
| if N > MAX_POINTS: |
| idx = np.linspace(0, N - 1, MAX_POINTS).astype(int) |
| actual, vis = actual[:, idx], vis[:, idx] |
| N = actual.shape[1] |
|
|
| sx, sy = SQUASH / W, SQUASH / H |
| x0_384 = actual[0] * np.array([sx, sy], np.float32) |
| pred_384 = predict(model, horizon, x0_384, device) |
| pred = pred_384 / np.array([sx, sy], np.float32) |
| L = pred.shape[0] |
|
|
| |
| tt = np.linspace(0, L - 1, Td, dtype=np.float32) |
| t0 = np.floor(tt).astype(int).clip(0, L - 1) |
| t1 = np.minimum(t0 + 1, L - 1) |
| w = (tt - t0)[:, None, None] |
| pred_al = (1 - w) * pred[t0] + w * pred[t1] |
|
|
| out_frames = [] |
| for ti, f in enumerate(frames): |
| cur = f.copy() |
| for pi in range(N): |
| if vis[ti, pi]: |
| x, y = actual[ti, pi] |
| cv2.circle(cur, (int(x), int(y)), 2, (0, 255, 0), -1) |
| for pi in range(N): |
| x, y = pred_al[ti, pi] |
| cv2.circle(cur, (int(x), int(y)), 2, (255, 0, 0), -1) |
| err = np.linalg.norm(actual[ti] - pred_al[ti], axis=-1) |
| m = vis[ti].astype(np.float32) |
| me = float((err * m).sum() / max(m.sum(), 1)) |
| cv2.putText(cur, f"t={ti} mean px err={me:5.1f}", (12, 30), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 0), 2) |
| cv2.putText(cur, "GT (CoTracker)", (12, H - 40), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) |
| cv2.putText(cur, "3PoinTr pred", (12, H - 14), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) |
| out_frames.append(cur) |
|
|
| os.makedirs(out_dir, exist_ok=True) |
| out = f"{out_dir}/pred_flow_{ep}.mp4" |
| imageio.mimwrite(out, out_frames, fps=15, macro_block_size=1) |
| err = np.linalg.norm(actual - pred_al, axis=-1) * vis |
| denom = vis.sum(1).clip(min=1) |
| me = (err.sum(1) / denom) |
| print(f"[{ep}] N={N} L={L} Td={Td} px err first={me[0]:.1f} mid={me[Td//2]:.1f} " |
| f"last={me[-1]:.1f} mean={me.mean():.1f} -> {out}") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ep", action="append", default=[]) |
| ap.add_argument("--all", action="store_true") |
| ap.add_argument("--ckpt", default=CKPT) |
| ap.add_argument("--out_dir", default="/home/hongyi/PWM/external/pred_flow_pushbox") |
| ap.add_argument("--device", default="cuda:0") |
| a = ap.parse_args() |
| eps = a.ep |
| if a.all: |
| import glob |
| eps = sorted(os.path.basename(p).replace("tableHF_", "") |
| for p in glob.glob(f"{TRACK_CACHE}/tableHF_ep*")) |
| if not eps: |
| eps = ["ep000"] |
| device = torch.device(a.device) |
| model, horizon = load_model(a.ckpt, device) |
| for ep in eps: |
| try: |
| render(ep, model, horizon, device, a.out_dir) |
| except Exception as e: |
| print(f"[{ep}] FAILED: {e}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
