PWM / PWM_code /reward_model /scripts /gen_pred_flow_openlaptop.py

Add files using upload-large-folder tool

dd253d4 verified 3 months ago

7.51 kB

#!/usr/bin/env python3
"""Open-Laptop point-flow overlay: GT CoTracker tracks vs 3-PoinTr pred (SINGLE object).
Tracks ONE object (the laptop lid). Model open_laptop_2d_128pts_384 is NOT
label-conditioned. GREEN=GT, RED=pred. Tracks are 384-space; frames come from the
EgoVLA-SIM hdf5 (observations/images/main, 384), drawn upscaled 2x. horizon inferred.
Usage: python gen_pred_flow_openlaptop.py --ep 0 --ep 2 --ep 6 --device cuda:0
"""
import argparse, os, sys
import cv2, imageio, h5py
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/open_laptop_2d_128pts_384_flow_backbone_final.pth"
TRACKS_ROOT = "/data0/hongyi/reward_model_data/3pointr_data_open_laptop_2d"   # episode_N/tracks_labeled.npz
HDF5 = "/data0/hongyi/EgoVLA_SIM/Open-Laptop"
IMG_KEY = "observations/images/main"
SQUASH = 384
UPSCALE = 2
MAX_POINTS = 128


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", 167))
    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, device):
    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}
    out = model(batch, mode="eval")
    pred = out["pred_pos"]
    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 load_frames(ep_idx):
    with h5py.File(f"{HDF5}/episode_{ep_idx}.hdf5", "r") as h:
        rgb = h[IMG_KEY][:]                      # (T,384,384,3) RGB
    fr = []
    for f in rgb:
        f = np.ascontiguousarray(f)
        if UPSCALE != 1:
            f = cv2.resize(f, (SQUASH * UPSCALE, SQUASH * UPSCALE), interpolation=cv2.INTER_LINEAR)
        fr.append(f)
    return fr


def render(ep_idx, model, horizon, device, out_dir, hold_tail=0):
    z = np.load(f"{TRACKS_ROOT}/episode_{ep_idx}/tracks_labeled.npz")
    tr = z["tracks"].astype(np.float32)                     # (T,N,2)@384
    vis = z["vis"] if "vis" in z.files else np.ones(tr.shape[:2], bool)
    N = tr.shape[1]
    if N > MAX_POINTS:
        idx = np.linspace(0, N - 1, MAX_POINTS).astype(int)
        tr, vis = tr[:, idx], vis[:, idx]; N = tr.shape[1]

    frames = load_frames(ep_idx)
    # match the hold-tail used in training: extend the display video + GT by repeating
    # the LAST frame/positions, so the model's held tail lines up 1:1 (else the 151-step
    # pred gets squished onto the 122-frame video and the hold isn't visible).
    if hold_tail > 0:
        frames = frames + [frames[-1].copy() for _ in range(hold_tail)]
        tr = np.concatenate([tr, np.repeat(tr[-1:], hold_tail, axis=0)], axis=0)
        vis = np.concatenate([vis, np.repeat(vis[-1:], hold_tail, axis=0)], axis=0)
    Td = min(len(frames), tr.shape[0])
    H, W = frames[0].shape[:2]
    s = float(UPSCALE)

    pred384 = predict(model, horizon, tr[0], device)
    pred = pred384 * s
    gt = tr * s
    L = pred.shape[0]
    # 1:1 demo-step mapping. pred[k] is the model's prediction for trajectory step k; this
    # demo was trained at length Td (tracked + hold). The model's horizon L is the MAX over
    # ALL demos, so pred[Td:L] is BEYOND this demo's supervised range (the flow extrapolates/
    # relaxes there) — resampling video frames onto [0, L-1] reads that region and looks like
    # the points "drop back". Take the first Td supervised steps, 1:1 with the Td frames.
    Td = min(Td, L)
    pred_al = pred[:Td]
    gt = gt[:Td]

    # vis collapses to ~0 (ignore-vis regime) → show ALL GT points (CoTracker tracks through occlusion)
    out_frames = []
    for ti in range(Td):
        cur = frames[ti].copy()
        for pi in range(N):
            x, y = gt[ti, pi]; cv2.circle(cur, (int(x), int(y)), 2, (0, 255, 0), -1)   # GT green
        for pi in range(N):
            x, y = pred_al[ti, pi]; cv2.circle(cur, (int(x), int(y)), 2, (255, 0, 0), -1)  # pred red
        err = np.linalg.norm(gt[ti] - pred_al[ti], axis=-1)
        me = float(err.mean()) / s
        cv2.putText(cur, f"ep{ep_idx}  t={ti}  mean err={me:4.1f}px@384", (10, 26),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
        cv2.putText(cur, "GT laptop lid", (10, H - 34), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        cv2.putText(cur, "3PoinTr pred", (10, H - 12), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
        out_frames.append(cur)

    os.makedirs(out_dir, exist_ok=True)
    out = f"{out_dir}/pred_flow_ep{ep_idx}.mp4"
    imageio.mimwrite(out, out_frames, fps=15, macro_block_size=1)
    err = np.linalg.norm(gt[:Td] - pred_al, axis=-1)
    me = err.mean(1) / s
    print(f"[ep{ep_idx}] N={N} L={L} Td={Td}  px err@384 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=[])      # episode indices, e.g. --ep 0 --ep 2
    ap.add_argument("--ckpt", default=CKPT)
    ap.add_argument("--out_dir", default="/home/hongyi/PWM/external/pred_flow_openlaptop")
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--hold_tail", type=int, default=0, help="extend display video+GT by H held frames")
    a = ap.parse_args()
    eps = a.ep or ["0"]
    device = torch.device(a.device)
    model, horizon = load_model(a.ckpt, device)
    for ep in eps:
        try:
            render(int(ep), model, horizon, device, a.out_dir, hold_tail=a.hold_tail)
        except Exception as e:
            import traceback; traceback.print_exc()
            print(f"[ep{ep}] FAILED: {e}")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Open-Laptop point-flow overlay: GT CoTracker tracks vs 3-PoinTr pred (SINGLE object).

	Tracks ONE object (the laptop lid). Model open_laptop_2d_128pts_384 is NOT
	label-conditioned. GREEN=GT, RED=pred. Tracks are 384-space; frames come from the
	EgoVLA-SIM hdf5 (observations/images/main, 384), drawn upscaled 2x. horizon inferred.

	Usage: python gen_pred_flow_openlaptop.py --ep 0 --ep 2 --ep 6 --device cuda:0
	"""
	import argparse, os, sys
	import cv2, imageio, h5py
	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/open_laptop_2d_128pts_384_flow_backbone_final.pth"
	TRACKS_ROOT = "/data0/hongyi/reward_model_data/3pointr_data_open_laptop_2d" # episode_N/tracks_labeled.npz
	HDF5 = "/data0/hongyi/EgoVLA_SIM/Open-Laptop"
	IMG_KEY = "observations/images/main"
	SQUASH = 384
	UPSCALE = 2
	MAX_POINTS = 128


	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", 167))
	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, device):
	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}
	out = model(batch, mode="eval")
	pred = out["pred_pos"]
	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 load_frames(ep_idx):
	with h5py.File(f"{HDF5}/episode_{ep_idx}.hdf5", "r") as h:
	rgb = h[IMG_KEY][:] # (T,384,384,3) RGB
	fr = []
	for f in rgb:
	f = np.ascontiguousarray(f)
	if UPSCALE != 1:
	f = cv2.resize(f, (SQUASH * UPSCALE, SQUASH * UPSCALE), interpolation=cv2.INTER_LINEAR)
	fr.append(f)
	return fr


	def render(ep_idx, model, horizon, device, out_dir, hold_tail=0):
	z = np.load(f"{TRACKS_ROOT}/episode_{ep_idx}/tracks_labeled.npz")
	tr = z["tracks"].astype(np.float32) # (T,N,2)@384
	vis = z["vis"] if "vis" in z.files else np.ones(tr.shape[:2], bool)
	N = tr.shape[1]
	if N > MAX_POINTS:
	idx = np.linspace(0, N - 1, MAX_POINTS).astype(int)
	tr, vis = tr[:, idx], vis[:, idx]; N = tr.shape[1]

	frames = load_frames(ep_idx)
	# match the hold-tail used in training: extend the display video + GT by repeating
	# the LAST frame/positions, so the model's held tail lines up 1:1 (else the 151-step
	# pred gets squished onto the 122-frame video and the hold isn't visible).
	if hold_tail > 0:
	frames = frames + [frames[-1].copy() for _ in range(hold_tail)]
	tr = np.concatenate([tr, np.repeat(tr[-1:], hold_tail, axis=0)], axis=0)
	vis = np.concatenate([vis, np.repeat(vis[-1:], hold_tail, axis=0)], axis=0)
	Td = min(len(frames), tr.shape[0])
	H, W = frames[0].shape[:2]
	s = float(UPSCALE)

	pred384 = predict(model, horizon, tr[0], device)
	pred = pred384 * s
	gt = tr * s
	L = pred.shape[0]
	# 1:1 demo-step mapping. pred[k] is the model's prediction for trajectory step k; this
	# demo was trained at length Td (tracked + hold). The model's horizon L is the MAX over
	# ALL demos, so pred[Td:L] is BEYOND this demo's supervised range (the flow extrapolates/
	# relaxes there) — resampling video frames onto [0, L-1] reads that region and looks like
	# the points "drop back". Take the first Td supervised steps, 1:1 with the Td frames.
	Td = min(Td, L)
	pred_al = pred[:Td]
	gt = gt[:Td]

	# vis collapses to ~0 (ignore-vis regime) → show ALL GT points (CoTracker tracks through occlusion)
	out_frames = []
	for ti in range(Td):
	cur = frames[ti].copy()
	for pi in range(N):
	x, y = gt[ti, pi]; cv2.circle(cur, (int(x), int(y)), 2, (0, 255, 0), -1) # GT green
	for pi in range(N):
	x, y = pred_al[ti, pi]; cv2.circle(cur, (int(x), int(y)), 2, (255, 0, 0), -1) # pred red
	err = np.linalg.norm(gt[ti] - pred_al[ti], axis=-1)
	me = float(err.mean()) / s
	cv2.putText(cur, f"ep{ep_idx} t={ti} mean err={me:4.1f}px@384", (10, 26),
	cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
	cv2.putText(cur, "GT laptop lid", (10, H - 34), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
	cv2.putText(cur, "3PoinTr pred", (10, H - 12), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
	out_frames.append(cur)

	os.makedirs(out_dir, exist_ok=True)
	out = f"{out_dir}/pred_flow_ep{ep_idx}.mp4"
	imageio.mimwrite(out, out_frames, fps=15, macro_block_size=1)
	err = np.linalg.norm(gt[:Td] - pred_al, axis=-1)
	me = err.mean(1) / s
	print(f"[ep{ep_idx}] N={N} L={L} Td={Td} px err@384 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=[]) # episode indices, e.g. --ep 0 --ep 2
	ap.add_argument("--ckpt", default=CKPT)
	ap.add_argument("--out_dir", default="/home/hongyi/PWM/external/pred_flow_openlaptop")
	ap.add_argument("--device", default="cuda:0")
	ap.add_argument("--hold_tail", type=int, default=0, help="extend display video+GT by H held frames")
	a = ap.parse_args()
	eps = a.ep or ["0"]
	device = torch.device(a.device)
	model, horizon = load_model(a.ckpt, device)
	for ep in eps:
	try:
	render(int(ep), model, horizon, device, a.out_dir, hold_tail=a.hold_tail)
	except Exception as e:
	import traceback; traceback.print_exc()
	print(f"[ep{ep}] FAILED: {e}")


	if __name__ == "__main__":
	main()