File size: 7,020 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 | #!/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()
|