#!/usr/bin/env python3 """Stack-Can point-flow overlay: GT CoTracker tracks vs label-conditioned 3-PoinTr pred. GREEN = GT moving-points (SAM3+CoTracker can+plate tracks from tracks_labeled.npz). RED = 3-PoinTr predicted future trajectory, rolled out from frame-0 GT points + their per-point labels (0=can, 1=plate) — the can+plate label-conditioned model. Tracks live in 384-squash space (training domain); we predict there and rescale the flow back to native 1280x720 to draw a crisp overlay on the original demo video. Usage: python gen_pred_flow_stackcan.py --ep ep000 [--ep ep007 ...] """ 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/stack_can_canplate_150pts_384_flow_backbone.pth" TRACKS_ROOT = "/data0/hongyi/reward_model_data/3pointr_data_stack_can_2d_canplate" # tracks_labeled.npz (384 space) DEMO_ROOT = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can" # native 1280x720 SQUASH = 384 MAX_POINTS = 150 # match training subsample 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", 252)) use_labels = any("point_label_embed" in k for k in state.keys()) 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, use_point_labels=use_labels, ).to(device) missing, unexpected = model.load_state_dict(state, strict=False) print(f"[predict] horizon={horizon} use_point_labels={use_labels} " f"missing={len(missing)} unexpected={len(unexpected)}") model.eval() return model, horizon @torch.no_grad() def predict(model, horizon, x0_px, labels, device): """x0_px (N,2)@384, labels (N,) -> (L,N,2)@384.""" 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) batch = {"x0": x0n, "num_steps": horizon, "task_ids": torch.zeros(1, dtype=torch.long, device=device), "scales": scales, "center": center, "point_labels": torch.from_numpy(labels.astype(np.int64)).to(device).unsqueeze(0)} out = model(batch, mode="eval") pred = out["pred_pos"] # (1,L,N,2) 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): z = np.load(f"{TRACKS_ROOT}/tableHF_{ep}/tracks_labeled.npz", allow_pickle=True) tr = z["tracks"].astype(np.float32) # (T,N,2)@384 lab = z["labels"].astype(np.int64) vis = z["vis"] if "vis" in z.files else np.ones(tr.shape[:2], bool) N = tr.shape[1] if N > MAX_POINTS: # even subsample (keeps label order) idx = np.linspace(0, N - 1, MAX_POINTS).astype(int) tr, vis, lab = tr[:, idx], vis[:, idx], lab[idx] N = tr.shape[1] cap = cv2.VideoCapture(f"{DEMO_ROOT}/tableHF_{ep}.mp4") frames = [] while True: ok, f = cap.read() if not ok: break frames.append(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) cap.release() Td = min(len(frames), tr.shape[0]) H, W = frames[0].shape[:2] sx, sy = W / SQUASH, H / SQUASH # 384 -> native pred384 = predict(model, horizon, tr[0], lab, device) # (L,N,2)@384 pred = pred384 * np.array([sx, sy], np.float32) # -> native gt = tr * np.array([sx, sy], np.float32) # -> native 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] # (Td,N,2) out_frames = [] for ti in range(Td): cur = frames[ti].copy() for pi in range(N): if vis[ti, pi]: x, y = gt[ti, pi]; cv2.circle(cur, (int(x), int(y)), 3, (0, 255, 0), -1) # GT green for pi in range(N): x, y = pred_al[ti, pi]; cv2.circle(cur, (int(x), int(y)), 3, (255, 0, 0), -1) # pred red err = np.linalg.norm(gt[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"{ep} t={ti} mean px err={me:5.1f}", (14, 36), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 0), 2) cv2.putText(cur, "GT (CoTracker)", (14, H - 46), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2) cv2.putText(cur, "3PoinTr pred", (14, H - 16), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (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(gt[:Td] - pred_al, axis=-1) * vis[:Td] me = err.sum(1) / vis[:Td].sum(1).clip(min=1) 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("--ckpt", default=CKPT) ap.add_argument("--out_dir", default="/home/hongyi/PWM/external/pred_flow_stackcan") ap.add_argument("--device", default="cuda:0") a = ap.parse_args() eps = a.ep or ["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()