| """Overlay 3-PoinTr object tracks (3D positions, projected via stored intrinsics) |
| onto the existing hand-action _pred.mp4 frames, producing _pred_handobject.mp4. |
| |
| The hand silhouettes (black) are preserved; object tracks are drawn only on the |
| white background. Using small colored dots avoids visually colliding with the |
| hand silhouettes the WM is already conditioned on. |
| """ |
| import argparse, os, json |
| import numpy as np |
| import torch |
| import cv2 |
| import imageio.v2 as imageio |
|
|
|
|
| def load_video(path): |
| cap = cv2.VideoCapture(path); frames = [] |
| while True: |
| r, f = cap.read() |
| if not r: break |
| frames.append(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) |
| cap.release() |
| return frames |
|
|
|
|
| def load_episode(ep_dir): |
| meta = json.load(open(os.path.join(ep_dir, "metadata.json"))) |
| T = meta["T"] |
| pos, vis = [], [] |
| for k in range(1, T + 1): |
| ts = os.path.join(ep_dir, f"timestep_{k}") |
| pos.append(torch.load(os.path.join(ts, f"{k}_pos.pth"), weights_only=False).numpy()) |
| vis.append(torch.load(os.path.join(ts, f"{k}_visibility.pth"), weights_only=False).numpy()) |
| pos = np.stack(pos, 0); vis = np.stack(vis, 0) |
| K = meta["intrinsics"]["camera"] |
| return pos, vis, (K["fx"], K["fy"], K["cx"], K["cy"]), meta |
|
|
|
|
| def project(p, fx, fy, cx, cy): |
| z = np.maximum(p[..., 2], 1e-6) |
| u = p[..., 0] / z * fx + cx |
| v = p[..., 1] / z * fy + cy |
| return np.stack([u, v], -1) |
|
|
|
|
| def is_white(pixel, thresh=235): |
| return (pixel >= thresh).all() |
|
|
|
|
| def overlay_object_tracks(pred_frames, pos_3d, vis, K, *, color="red", |
| dot_radius=3, stride=1, draw_only_on_white=True, |
| white_thresh=230, trail_len=0): |
| """Draw colored dots on the pred frames at projected object positions. |
| |
| `draw_only_on_white=True` ensures hand silhouettes (black) are never overwritten, |
| so the WM can still see the same hand instruction. |
| """ |
| fx, fy, cx, cy = K |
| Tp = len(pred_frames) |
| Tt = pos_3d.shape[0] |
| T = min(Tp, Tt) |
| H, W = pred_frames[0].shape[:2] |
|
|
| color_map = { |
| "red": (220, 30, 30), |
| "blue": (30, 120, 255), |
| "green": (40, 200, 40), |
| "magenta": (220, 30, 200), |
| } |
| col_rgb = color_map[color] |
|
|
| out = [] |
| proj = project(pos_3d, fx, fy, cx, cy) |
| for t in range(T): |
| f = pred_frames[t].copy() |
| |
| for n in range(0, pos_3d.shape[1], stride): |
| if not vis[t, n]: |
| continue |
| x, y = proj[t, n] |
| if not (0 <= x < W and 0 <= y < H): |
| continue |
| xi, yi = int(round(x)), int(round(y)) |
| if draw_only_on_white and not is_white(f[yi, xi], white_thresh): |
| |
| continue |
| cv2.circle(f, (xi, yi), dot_radius, col_rgb, -1, cv2.LINE_AA) |
| |
| if trail_len > 0 and t > 0: |
| t0 = max(0, t - trail_len) |
| for n in range(0, pos_3d.shape[1], stride): |
| for ti in range(t0 + 1, t + 1): |
| if not (vis[ti, n] and vis[ti - 1, n]): |
| continue |
| x1, y1 = proj[ti - 1, n] |
| x2, y2 = proj[ti, n] |
| if not (0 <= x1 < W and 0 <= y1 < H and 0 <= x2 < W and 0 <= y2 < H): |
| continue |
| cv2.line(f, (int(x1), int(y1)), (int(x2), int(y2)), col_rgb, 1, cv2.LINE_AA) |
| out.append(f) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--pred_video", required=True) |
| ap.add_argument("--ep_dir", required=True, |
| help="3-PoinTr episode dir, e.g. .../no_actions/stack_can/0001") |
| ap.add_argument("--out_video", required=True) |
| ap.add_argument("--color", default="red", choices=["red", "blue", "green", "magenta"]) |
| ap.add_argument("--stride", type=int, default=1) |
| ap.add_argument("--dot_radius", type=int, default=3) |
| ap.add_argument("--trail_len", type=int, default=0) |
| ap.add_argument("--no_white_only", action="store_true", |
| help="Allow drawing on hand silhouette too (default protects hands).") |
| args = ap.parse_args() |
|
|
| print(f"[in] pred : {args.pred_video}") |
| print(f"[in] ep_dir: {args.ep_dir}") |
| pred_frames = load_video(args.pred_video) |
| pos, vis, K, meta = load_episode(args.ep_dir) |
| print(f" pred T={len(pred_frames)} tracks T={pos.shape[0]} N={pos.shape[1]}") |
| print(f" intrinsics: fx={K[0]:.1f} fy={K[1]:.1f} cx={K[2]:.1f} cy={K[3]:.1f}") |
|
|
| out = overlay_object_tracks( |
| pred_frames, pos, vis, K, |
| color=args.color, dot_radius=args.dot_radius, stride=args.stride, |
| draw_only_on_white=(not args.no_white_only), |
| trail_len=args.trail_len, |
| ) |
| imageio.mimwrite(args.out_video, out, fps=15, macro_block_size=1) |
| print(f"[out] {args.out_video} ({len(out)} frames)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
