| """Load trained 3-PoinTr backbone, predict trajectories on a val episode, |
| project predictions back to image, overlay GT and predicted points side-by-side.""" |
| import os, sys, json |
| sys.path.insert(0, "/home/hongyi/PWM/reward_model/3PoinTr_repo") |
|
|
| import numpy as np |
| import torch |
| import imageio.v2 as imageio |
| import cv2 |
|
|
| from util.training_utils.arch import TrajectoryModel |
| from util.training_utils.train_util import ( |
| build_common_arg_parser, prepare_datasets_and_loaders, normalize_and_augment, |
| extract_diffusion_kwargs, |
| ) |
|
|
|
|
| DATA_ROOT = "/home/hongyi/PWM/external/3pointr_data" |
| SPLIT_FILE = os.path.join(DATA_ROOT, "stack_can_flow.json") |
| CKPT = "/home/hongyi/PWM/reward_model/3pointr_models/stack_can_flowonly_512pts_flow_backbone.pth" |
| EP_DIR = "/home/hongyi/PWM/external/3pointr_data/no_actions/stack_can/0036" |
| OUT_VIDEO = "/home/hongyi/PWM/reward_model/3pointr_models/val_ep0036_pred_vs_gt.mp4" |
|
|
|
|
| def load_ep(ep_dir): |
| meta = json.load(open(os.path.join(ep_dir, "metadata.json"))) |
| T = meta["T"] |
| pos, vis, rgbs = [], [], [] |
| 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")).numpy()) |
| vis.append(torch.load(os.path.join(ts, f"{k}_visibility.pth")).numpy()) |
| rgbs.append(imageio.imread(os.path.join(ts, f"{k}_rgb.png"))) |
| pos = np.stack(pos, 0); vis = np.stack(vis, 0) |
| K = meta["intrinsics"]["camera"] |
| return rgbs, 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 main(): |
| parser = build_common_arg_parser("predict") |
| args = parser.parse_args([ |
| "--dataset", DATA_ROOT + "/", "--split_file", SPLIT_FILE, |
| "--subsample", "512", "--batch_size", "1", |
| "--num_queries", "10", |
| "--arch_dim", "256", "--arch_depth", "2", "--arch_heads", "4", |
| "--arch_query_depth", "1", "--arch_query_heads", "1", |
| "--flow_data_only", "--tasks", "stack_can", |
| ]) |
|
|
| arts = prepare_datasets_and_loaders(args) |
| traj_len = arts.train_set.traj_len |
| num_tasks = len(arts.tasks) |
| print(f"traj_len={traj_len}, num_tasks={num_tasks}") |
|
|
| diff_kwargs = extract_diffusion_kwargs(args) |
| diff_kwargs["n_action_steps"] = traj_len - 1 |
| model = TrajectoryModel( |
| dim=args.arch_dim, depth=args.arch_depth, num_heads=args.arch_heads, |
| query_depth=args.arch_query_depth, query_heads=args.arch_query_heads, |
| horizon=traj_len - 1, traj_len=traj_len, |
| use_task_conditioning=False, num_tasks=num_tasks, |
| decoder_only_modules=("decoder_only.point_flow_head",), |
| **diff_kwargs, |
| ).cuda().eval() |
| sd = torch.load(CKPT, map_location="cuda") |
| if isinstance(sd, dict) and "model" in sd: sd = sd["model"] |
| if isinstance(sd, dict) and "state_dict" in sd: sd = sd["state_dict"] |
| missing, unexpected = model.load_state_dict(sd, strict=False) |
| print(f"loaded ckpt: missing={len(missing)} unexpected={len(unexpected)}") |
|
|
| rgbs, pos_gt, vis_gt, (fx, fy, cx, cy), meta = load_ep(EP_DIR) |
| T_ep = pos_gt.shape[0]; N = pos_gt.shape[1]; H, W = rgbs[0].shape[:2] |
| print(f"ep T={T_ep}, N={N}") |
|
|
| |
| x0 = torch.from_numpy(pos_gt[0]).float().unsqueeze(0) |
| cum = torch.from_numpy(pos_gt - pos_gt[0:1]).float() |
| |
| L = traj_len - 1 |
| cum_pad = torch.zeros(L, N, 3) |
| pmask = torch.zeros(L, N, dtype=torch.bool) |
| smask = torch.zeros(L, dtype=torch.bool) |
| nsteps = min(L, T_ep - 1) |
| cum_pad[:nsteps] = torch.from_numpy(pos_gt[1:1+nsteps] - pos_gt[0:1]) |
| pmask[:nsteps] = torch.from_numpy(vis_gt[1:1+nsteps]).bool() |
| smask[:nsteps] = True |
| cum = cum_pad.unsqueeze(0); point_mask = pmask.unsqueeze(0); step_mask = smask.unsqueeze(0) |
| action_mask = torch.zeros(1, dtype=torch.bool) |
|
|
| norm = normalize_and_augment(args, x0, cum, None, None, None, |
| step_mask, action_mask, point_mask, "cuda", |
| apply_noise=False) |
| feed = {"x0": norm["x0n"], "step_mask": norm["step_mask"], |
| "point_mask": norm["point_mask"], "num_steps": L, |
| "flow_loss_w": 0.0, "cum_target": norm["cum_gtn"]} |
| with torch.no_grad(): |
| out = model(feed, mode="eval") |
| pred_pos_norm = out["pred_pos"][0].cpu().numpy() |
|
|
| |
| scales = norm["scales"][0].cpu().numpy() |
| center = (x0[0].numpy()).mean(axis=0) |
| |
| |
| |
| |
| |
| |
| pred_pos_world = pred_pos_norm * scales + center |
| print(f"pred range world: x={pred_pos_world[:,:,0].min():.3f}..{pred_pos_world[:,:,0].max():.3f} " |
| f"z={pred_pos_world[:,:,2].min():.3f}..{pred_pos_world[:,:,2].max():.3f}") |
| print(f"GT range world: x={pos_gt[:,:,0].min():.3f}..{pos_gt[:,:,0].max():.3f} " |
| f"z={pos_gt[:,:,2].min():.3f}..{pos_gt[:,:,2].max():.3f}") |
|
|
| |
| pred_uv = project(pred_pos_world, fx, fy, cx, cy) |
| gt_uv = project(pos_gt[1:1+nsteps], fx, fy, cx, cy) |
| init_uv = project(pos_gt[0:1], fx, fy, cx, cy)[0] |
|
|
| out_frames = [] |
| for t in range(T_ep): |
| f = rgbs[t].copy() |
| |
| for n in range(0, N, 4): |
| x, y = init_uv[n] |
| if 0 <= x < W and 0 <= y < H: |
| cv2.circle(f, (int(x), int(y)), 3, (50, 100, 255), -1, cv2.LINE_AA) |
| |
| if t > 0 and (t - 1) < gt_uv.shape[0]: |
| for n in range(0, N, 4): |
| if not vis_gt[t, n]: continue |
| x, y = gt_uv[t-1, n] |
| if 0 <= x < W and 0 <= y < H: |
| cv2.circle(f, (int(x), int(y)), 3, (0, 255, 0), -1, cv2.LINE_AA) |
| |
| if t > 0 and (t - 1) < pred_uv.shape[0]: |
| for n in range(0, N, 4): |
| x, y = pred_uv[t-1, n] |
| if 0 <= x < W and 0 <= y < H: |
| cv2.circle(f, (int(x), int(y)), 3, (255, 50, 50), -1, cv2.LINE_AA) |
| cv2.putText(f, f"t={t} blue=init green=GT red=PRED", (16, 32), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2, cv2.LINE_AA) |
| out_frames.append(f) |
|
|
| imageio.mimwrite(OUT_VIDEO, out_frames, fps=15, macro_block_size=1) |
| print(f"saved {OUT_VIDEO}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
