"""Evaluate trained 3-PoinTr backbone on val + 1 train episode: • Per-episode EPE (in mm) on visible points • Per-timestep EPE curve (in mm) • Overlay video: blue=init points, green=GT trail, red=pred trail """ import os, sys, json, glob 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" OUT_DIR = "/home/hongyi/PWM/reward_model/3pointr_models/eval" os.makedirs(OUT_DIR, exist_ok=True) 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"), weights_only=False).numpy()) vis.append(torch.load(os.path.join(ts, f"{k}_visibility.pth"), weights_only=False).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 predict_one(model, args, pos_gt, vis_gt, traj_len): T_ep, N, _ = pos_gt.shape L = traj_len - 1 nsteps = min(L, T_ep - 1) x0 = torch.from_numpy(pos_gt[0]).float().unsqueeze(0) cum_pad = torch.zeros(L, N, 3) pmask = torch.zeros(L, N, dtype=torch.bool) smask = torch.zeros(L, dtype=torch.bool) 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_norm = out["pred_pos"][0].cpu().numpy() # (L, N, 3) normalized # Inverse of normalize_and_augment using the single-scalar radius and 3D center. scales = norm["scales"][0].cpu().numpy().item() # scalar center = norm["center"][0].cpu().numpy() # (3,) pred_world = pred_norm * scales + center return pred_world[:nsteps], nsteps def _draw_panel(frame, pos_init_uv, gt_uv_seq, pred_uv_seq, vis_gt, t, show_gt, show_pred, trail_len=15, stride=4): H, W = frame.shape[:2] f = frame.copy() # init (blue) — always for n in range(0, pos_init_uv.shape[0], stride): x, y = pos_init_uv[n] if 0 <= x < W and 0 <= y < H: cv2.circle(f, (int(x), int(y)), 2, (50, 100, 255), -1, cv2.LINE_AA) if t > 0: t0 = max(1, t - trail_len) if show_gt: for n in range(0, gt_uv_seq.shape[1], stride): for ti in range(t0, min(t + 1, gt_uv_seq.shape[0] + 1)): if ti - 1 < 0 or ti - 1 >= gt_uv_seq.shape[0]: continue if not vis_gt[ti, n]: continue x, y = gt_uv_seq[ti - 1, n] if 0 <= x < W and 0 <= y < H: cv2.circle(f, (int(x), int(y)), 2, (0, 255, 0), -1, cv2.LINE_AA) if show_pred: for n in range(0, pred_uv_seq.shape[1], stride): for ti in range(t0, min(t + 1, pred_uv_seq.shape[0] + 1)): if ti - 1 < 0 or ti - 1 >= pred_uv_seq.shape[0]: continue x, y = pred_uv_seq[ti - 1, n] if 0 <= x < W and 0 <= y < H: cv2.circle(f, (int(x), int(y)), 2, (255, 50, 50), -1, cv2.LINE_AA) return f def overlay_video(rgbs, pos_init_uv, gt_uv_seq, pred_uv_seq, vis_gt, out_path, trail_len=15, stride=4): """3-panel layout: [blue+red | blue+green | blue+green+red].""" H, W = rgbs[0].shape[:2] T = len(rgbs) out = [] for t in range(T): left = _draw_panel(rgbs[t], pos_init_uv, gt_uv_seq, pred_uv_seq, vis_gt, t, show_gt=False, show_pred=True, trail_len=trail_len, stride=stride) cv2.putText(left, f"PRED only (blue=init, red=pred) t={t}", (16, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 255), 2, cv2.LINE_AA) mid = _draw_panel(rgbs[t], pos_init_uv, gt_uv_seq, pred_uv_seq, vis_gt, t, show_gt=True, show_pred=False, trail_len=trail_len, stride=stride) cv2.putText(mid, f"GT only (blue=init, green=GT) t={t}", (16, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 255), 2, cv2.LINE_AA) right = _draw_panel(rgbs[t], pos_init_uv, gt_uv_seq, pred_uv_seq, vis_gt, t, show_gt=True, show_pred=True, trail_len=trail_len, stride=stride) cv2.putText(right, f"BOTH (blue=init, green=GT, red=pred) t={t}", (16, 32), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 255), 2, cv2.LINE_AA) panel = np.concatenate([left, mid, right], axis=1) out.append(panel) imageio.mimwrite(out_path, out, fps=15, macro_block_size=1) def main(): parser = build_common_arg_parser("eval") 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 print(f"traj_len={traj_len}") 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=len(arts.tasks), decoder_only_modules=("decoder_only.point_flow_head",), **diff_kwargs, ).cuda().eval() raw = torch.load(CKPT, map_location="cuda", weights_only=False) sd = raw["model_state_dict"] if isinstance(raw, dict) and "model_state_dict" in raw else raw missing, unexpected = model.load_state_dict(sd, strict=False) print(f"ckpt loaded. missing={len(missing)} unexpected={len(unexpected)}") if "best_flow_epe_mm" in raw: print(f" ckpt's reported best val EPE: {raw['best_flow_epe_mm']:.4f} mm at epoch {raw.get('epoch', '?')}") # Pick: 3 val + 1 train for sanity split = json.load(open(SPLIT_FILE)) eps_to_eval = [ ("train_first", split["flows_train"][0]), ("val_0035", split["flows_val"][0]), ("val_0036", split["flows_val"][1]), ("val_0037", split["flows_val"][2]), ] summary = {} for tag, ep_dir in eps_to_eval: rgbs, pos_gt, vis_gt, (fx, fy, cx, cy), meta = load_ep(ep_dir) T_ep, N = pos_gt.shape[:2] pred_world, nsteps = predict_one(model, args, pos_gt, vis_gt, traj_len) # EPE per timestep (visible only), in mm gt_seq = pos_gt[1:1 + nsteps] # (nsteps, N, 3) v_seq = vis_gt[1:1 + nsteps] # (nsteps, N) diff = np.linalg.norm(pred_world - gt_seq, axis=-1) # (nsteps, N) mask_count = v_seq.sum(axis=-1).clip(min=1) epe_per_t = ((diff * v_seq).sum(axis=-1) / mask_count) * 1000 # mm epe_overall = ((diff * v_seq).sum() / max(1, v_seq.sum())) * 1000 summary[tag] = { "ep_dir": os.path.basename(ep_dir), "T": T_ep, "N": N, "epe_overall_mm": float(epe_overall), "epe_t1_mm": float(epe_per_t[0]), "epe_tmid_mm": float(epe_per_t[len(epe_per_t) // 2]), "epe_tlast_mm": float(epe_per_t[-1]), } print(f"{tag}: T={T_ep}, N={N}, EPE overall={epe_overall:.2f} mm, " f"t=1 {epe_per_t[0]:.2f}, t=mid {epe_per_t[len(epe_per_t)//2]:.2f}, " f"t=last {epe_per_t[-1]:.2f}") # Build overlay video gt_uv = project(gt_seq, fx, fy, cx, cy) pred_uv = project(pred_world, fx, fy, cx, cy) # Pad vis_gt[0] for indexing parity vis_full = vis_gt init_uv = project(pos_gt[0:1], fx, fy, cx, cy)[0] out_video = os.path.join(OUT_DIR, f"{tag}_pred_vs_gt.mp4") overlay_video(rgbs, init_uv, gt_uv, pred_uv, vis_full, out_video) print(f" saved {out_video}") # EPE curve plot (text-friendly) np.savetxt(os.path.join(OUT_DIR, f"{tag}_epe_per_t_mm.txt"), epe_per_t, fmt="%.3f") with open(os.path.join(OUT_DIR, "summary.json"), "w") as f: json.dump(summary, f, indent=2) print(f"\nsummary -> {OUT_DIR}/summary.json") if __name__ == "__main__": main()