Blame for PWM_code/reward_model/scripts/3pointr_eval.py

"""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()