File size: 7,100 Bytes
dd253d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | #!/usr/bin/env python3
"""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 # match the model's "150pts" 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) # out = L*point_dim
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) # (1,N,2)
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"] # (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):
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) # (T,N,2) native
vis = z["vis"] if "vis" in z.files else np.ones(actual.shape[:2], bool)
N = actual.shape[1]
if N > MAX_POINTS: # even subsample to ~150
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) # frame-0 GT -> 384 space
pred_384 = predict(model, horizon, x0_384, device) # (L,N,2) 384
pred = pred_384 / np.array([sx, sy], np.float32) # back to native
L = pred.shape[0]
# resample predicted (1..L) onto demo timeline (align pred step 0 -> demo frame 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, 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) # GT green
for pi in range(N):
x, y = pred_al[ti, pi]
cv2.circle(cur, (int(x), int(y)), 2, (255, 0, 0), -1) # pred red
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()
|