| |
| |
| """Render every annotated hand's predicted MANO mesh on the original image. |
| |
| Bbox annotations are a JSON of the form:: |
| |
| { |
| "image1.jpg": [ |
| {"bbox": [x1, y1, x2, y2], "is_right": 1}, |
| {"bbox": [x1, y1, x2, y2], "is_right": 0} |
| ], |
| ... |
| } |
| |
| Mirrors the upstream HaMeR demo: |
| |
| * Each hand bbox is processed independently (HaMeR is single-hand). |
| * For ``is_right=0`` (left hand) the crop is mirrored horizontally |
| before the model, then the predicted ``cam.x`` and ``verts.x`` are |
| flipped back — see ``hamer/demo.py`` lines 136-138 / 178-179. |
| * Mesh is projected onto the original frame via ``cam_crop_to_full`` |
| with ``focal = 5000 / 256 · max(W, H)`` (no spurious x-flip — that |
| only exists in upstream because ``pyrender`` places the camera at |
| ``cam_t``; here we translate the mesh in front of a fixed camera). |
| |
| Run from the repo root: |
| |
| python scripts/visualize.py \ |
| --ckpt ... --mean ... --mano ... \ |
| --images media --bboxes media/bboxes.json --out media |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
| _REPO_ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(_REPO_ROOT)) |
|
|
| from hamer import Hamer, load_hamer_checkpoint |
| from hamer.preprocess import preprocess |
| from hamer.mano import MANO |
| from hamer.rendering import overlay_mesh, project_mesh |
|
|
|
|
| def parse_args(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--ckpt", type=Path, required=True) |
| p.add_argument("--mean", type=Path, required=True) |
| p.add_argument("--mano", type=Path, required=True, help="MANO_RIGHT.pkl") |
| p.add_argument("--images", type=Path, required=True, help="dir of images") |
| p.add_argument("--bboxes", type=Path, required=True, |
| help="JSON: {image_name: [{bbox, is_right}, ...]}") |
| p.add_argument("--out", type=Path, default=Path("media")) |
| p.add_argument("--device-id", type=int, default=int(os.environ.get("HAMER_DEVICE", "0"))) |
| return p.parse_args() |
|
|
|
|
| def split_output(out: torch.Tensor): |
| rotmats = out[:, :144].reshape(-1, 16, 3, 3) |
| betas = out[:, 144:154] |
| cam = out[:, 154:157] |
| return rotmats, betas, cam |
|
|
|
|
| def normalize_entry(e): |
| """Accept either ``[x1,y1,x2,y2]`` (legacy single-hand) or |
| ``{bbox: [...], is_right: 0/1}``. Returns dict with ``bbox`` + |
| ``is_right``.""" |
| if isinstance(e, dict): |
| return {"bbox": tuple(e["bbox"]), "is_right": int(e.get("is_right", 1))} |
| return {"bbox": tuple(e), "is_right": 1} |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| args.out.mkdir(parents=True, exist_ok=True) |
| raw_bboxes = json.loads(args.bboxes.read_text()) |
| bboxes = { |
| name: [normalize_entry(e) for e in (entries if isinstance(entries, list) else [entries])] |
| for name, entries in raw_bboxes.items() |
| } |
|
|
| print("Loading reference + checkpoint…") |
| torch.manual_seed(0) |
| ref = Hamer(); ref.eval() |
| load_hamer_checkpoint(ref, args.ckpt, args.mean) |
| print("Loading MANO model…") |
| mano = MANO(args.mano) |
|
|
| os.environ["HAMER_USE_TT"] = "1" |
| os.environ["HAMER_FORCE_TT"] = "1" |
| from hamer import build_from_reference |
| print("Opening tt-nn device…") |
| tt_model = build_from_reference(ref, device_id=args.device_id) |
|
|
| |
| REF_COLORS = [(0.20, 0.85, 1.00), (1.00, 0.85, 0.20), (0.30, 1.00, 0.50), (1.00, 0.50, 0.85)] |
| TT_COLORS = [(1.00, 0.55, 0.20), (0.55, 0.85, 1.00), (1.00, 0.20, 0.40), (0.40, 0.95, 0.65)] |
|
|
| for name, hands in bboxes.items(): |
| img_path = args.images / name |
| if not img_path.exists(): |
| print(f" skip {name} — file not found") |
| continue |
| original = np.asarray(Image.open(img_path).convert("RGB")) |
| h, w = original.shape[:2] |
|
|
| ref_projs = [] |
| tt_projs = [] |
| for hand in hands: |
| x, bbox_meta = preprocess(img_path, bbox_xyxy=hand["bbox"], is_right=bool(hand["is_right"])) |
| cx, cy, bbox_size = bbox_meta |
| with torch.no_grad(): |
| ref_out = ref(x).detach() |
| tt_out = tt_model(x).detach() |
| rot_ref, betas_ref, cam_ref = split_output(ref_out) |
| rot_tt, betas_tt, cam_tt = split_output(tt_out) |
| with torch.no_grad(): |
| verts_ref = mano(rot_ref, betas_ref)[0] |
| verts_tt = mano(rot_tt, betas_tt)[0] |
| ref_projs.append(project_mesh(verts_ref, cam_ref[0], (cx, cy), bbox_size, w, h, is_right=bool(hand["is_right"]))) |
| tt_projs.append(project_mesh(verts_tt, cam_tt[0], (cx, cy), bbox_size, w, h, is_right=bool(hand["is_right"]))) |
|
|
| fig, axes = plt.subplots(1, 3, figsize=(3 * w / 90, h / 90)) |
| axes[0].imshow(original) |
| axes[0].set_title(f"input — {name}", fontsize=11) |
| axes[0].axis("off") |
|
|
| |
| axes[1].imshow(original) |
| axes[1].set_xlim(0, w); axes[1].set_ylim(h, 0) |
| axes[1].set_aspect("equal"); axes[1].axis("off") |
| axes[1].set_title("torch CPU reference meshes", fontsize=11) |
| for i, (hand, (v2d, vz)) in enumerate(zip(hands, ref_projs)): |
| overlay_mesh(axes[1], original, v2d, vz, mano.faces, color=REF_COLORS[i % len(REF_COLORS)], is_right=bool(hand["is_right"])) |
|
|
| |
| axes[2].imshow(original) |
| axes[2].set_xlim(0, w); axes[2].set_ylim(h, 0) |
| axes[2].set_aspect("equal"); axes[2].axis("off") |
| axes[2].set_title("tt-nn port meshes", fontsize=11) |
| for i, (hand, (v2d, vz)) in enumerate(zip(hands, tt_projs)): |
| overlay_mesh(axes[2], original, v2d, vz, mano.faces, color=TT_COLORS[i % len(TT_COLORS)], is_right=bool(hand["is_right"])) |
|
|
| plt.tight_layout(pad=0.5) |
| out_path = args.out / f"{Path(name).stem}_result.png" |
| plt.savefig(out_path, dpi=110, bbox_inches="tight") |
| plt.close(fig) |
| print(f" → {out_path}") |
|
|
| if hasattr(tt_model, "close"): |
| tt_model.close() |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
