| |
| """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" |
| DEMO_ROOT = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can" |
| 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", 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"] |
| 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) |
| 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: |
| 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 |
|
|
| pred384 = predict(model, horizon, tr[0], lab, device) |
| pred = pred384 * np.array([sx, sy], np.float32) |
| gt = tr * 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 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) |
| for pi in range(N): |
| x, y = pred_al[ti, pi]; cv2.circle(cur, (int(x), int(y)), 3, (255, 0, 0), -1) |
| 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() |
|
|
