| |
| |
| """Python API: HaMeR hand mesh recovery on one Tenstorrent Blackhole chip. |
| |
| from hamer import HamerModel |
| |
| with HamerModel.from_pretrained(device_id=0) as model: |
| hand = model("media/sample1.jpg", bbox=[20, 170, 295, 260], is_right=True) |
| print(hand.rotmats.shape, hand.betas, hand.cam_t_full) |
| |
| ``from_pretrained`` loads the released checkpoint from the Hugging Face cache, opens the chip the way the published |
| numbers were measured (1x1 mesh, ETH dispatch, 12x10 compute grid, the ``tt-model.yaml`` serve configuration), uploads |
| the weights, captures one ttnn trace and checks the device output against the torch CPU reference. Every call then |
| runs the same steps as ``POST /predict`` of the HTTP server: host crop + resize, one trace replay on the chip, host |
| post-processing (and the MANO mesh when ``MANO_RIGHT.pkl`` is available). The outputs are the same numbers. |
| |
| Nothing in this module imports ``ttnn`` at import time; the device stack loads inside ``from_pretrained``. |
| Full reference: ``PYTHON.md`` at the repository root. |
| """ |
| from __future__ import annotations |
|
|
| import functools |
| import gc |
| import os |
| import sys |
| import tempfile |
| import threading |
| import time |
| import warnings |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
| from .preprocess import ( |
| IMAGE_HEIGHT, IMAGE_WIDTH, W_SLICE, center_square_bbox, hamer_bbox_size, normalize_raw, preprocess_pil, |
| ) |
| from .rendering import FOCAL_LENGTH, MODEL_IMAGE_SIZE, cam_crop_to_full, project_mesh |
| from .served_config import ( |
| CKPT_FILE, MANO_FILE, MEAN_FILE, SERVED_DISPATCH, SERVED_ENV, SERVED_GRID, SERVED_WORKER_L1_SIZE, |
| TRACE_REGION_BYTES, WEIGHTS_REPO, WEIGHTS_REVISION, |
| ) |
|
|
| __all__ = ["HamerModel", "HandPrediction", "open_device", "load_reference", "apply_served_config", |
| "DEFAULT_WARMUP_VARIANTS"] |
|
|
| ImageLike = Union[str, os.PathLike, Image.Image, np.ndarray, torch.Tensor] |
| BBox = Sequence[float] |
|
|
| MIN_CROP_PIXELS = 2.0 |
| GRID_OVERRIDE_ENV = "TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE" |
|
|
|
|
| def _log(verbose: bool, msg: str) -> None: |
| if verbose: |
| print(f"[hamer] {msg}", flush=True) |
|
|
|
|
| |
| |
| QUIET_LOG_ENV = {"TT_LOGGER_LEVEL": "error", "LOGURU_LEVEL": "WARNING"} |
|
|
|
|
| def _quiet_logs(verbose: bool) -> Dict[str, str]: |
| """``verbose=False``: errors only from tt-metal / ttnn (when ``ttnn`` is not imported yet) and no ``[hamer]`` boot |
| lines. ``verbose=True``: the logging configuration is not changed. Returns the variables that were set.""" |
| from ._verbosity import set_quiet |
|
|
| set_quiet(not verbose) |
| done: Dict[str, str] = {} |
| if verbose or "ttnn" in sys.modules: |
| return done |
| for k, v in QUIET_LOG_ENV.items(): |
| if k not in os.environ: |
| os.environ[k] = v |
| done[k] = v |
| return done |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| WARMUP_FRAME_SIZE = (1920, 1080) |
| WARMUP_DEFAULTS: Dict[str, Dict[str, Any]] = { |
| "frame": {"input": "path", "fmt": "jpeg", "hands": 1, "is_right": True, "bbox": "box", |
| "frame_size": WARMUP_FRAME_SIZE}, |
| "crops": {"crop_dtype": "uint8", "crop_width": 256, "batch": 2}, |
| "run_crop": {"crop_dtype": "uint8"}, |
| } |
| WARMUP_CHOICES: Dict[str, Tuple[Any, ...]] = { |
| "input": ("path", "pil", "array", "tensor"), "fmt": ("jpeg", "png"), "bbox": ("box", "center"), |
| "crop_dtype": ("uint8", "float"), "crop_width": (256, IMAGE_WIDTH), "is_right": (True, False, "mixed"), |
| } |
| |
| |
| DEFAULT_WARMUP_VARIANTS: Tuple[Dict[str, Any], ...] = ( |
| {"call": "run_crop", "crop_dtype": "float"}, |
| {"call": "run_crop", "crop_dtype": "uint8"}, |
| {"call": "crops", "crop_dtype": "float", "crop_width": 256}, |
| {"call": "crops", "crop_dtype": "uint8", "crop_width": 256}, |
| {"call": "frame", "input": "tensor"}, |
| {"call": "frame", "input": "array"}, |
| {"call": "frame", "input": "pil", "bbox": "center"}, |
| {"call": "frame", "input": "path", "fmt": "png", "hands": 2, "frame_size": (640, 480)}, |
| {"call": "frame", "input": "path", "fmt": "jpeg", "is_right": False}, |
| {"call": "frame", "input": "path", "fmt": "jpeg", "is_right": True}, |
| ) |
|
|
|
|
| def normalize_warmup_variant(call: str = "frame", **variant: Any) -> Tuple[Tuple[str, Any], ...]: |
| """One warm-up variant as a sorted ``((key, value), ...)`` tuple with the defaults filled in (the idempotency key). |
| ``call``: ``"frame"`` (``model(image, bbox)``), ``"crops"`` (:meth:`HamerModel.forward_crops`) or ``"run_crop"``. |
| Raises ``ValueError`` for an unknown key or value.""" |
| if call not in WARMUP_DEFAULTS: |
| raise ValueError(f"warm-up call must be one of {sorted(WARMUP_DEFAULTS)}, got {call!r}") |
| allowed = WARMUP_DEFAULTS[call] |
| bad = sorted(set(variant) - set(allowed)) |
| if bad: |
| raise ValueError(f"unknown warm-up keys for call={call!r}: {bad} (allowed: {sorted(allowed)})") |
| v = dict(allowed, **variant) |
| if call == "frame": |
| if "is_right" not in variant and int(v["hands"]) > 1: |
| v["is_right"] = "mixed" |
| v["hands"] = int(v["hands"]) |
| if v["hands"] < 1: |
| raise ValueError(f"hands must be >= 1, got {v['hands']}") |
| if v["hands"] == 1 and v["is_right"] == "mixed": |
| raise ValueError("is_right='mixed' needs hands >= 2") |
| fs = tuple(int(x) for x in v["frame_size"]) |
| if len(fs) != 2 or min(fs) < 16 or fs[0] * fs[1] > 50_000_000: |
| raise ValueError(f"frame_size must be (W, H) with sides >= 16 and at most 50 MP, got {v['frame_size']}") |
| v["frame_size"] = fs |
| if v["input"] != "path": |
| v.pop("fmt") |
| if call == "crops": |
| v["batch"] = int(v["batch"]) |
| if v["batch"] < 1: |
| raise ValueError(f"batch must be >= 1, got {v['batch']}") |
| for k, val in v.items(): |
| if k in WARMUP_CHOICES and val not in WARMUP_CHOICES[k]: |
| raise ValueError(f"warm-up {k}={val!r} is not one of {WARMUP_CHOICES[k]}") |
| return (("call", call),) + tuple(sorted(v.items())) |
|
|
|
|
| def _warmup_spec(spec: Any) -> List[Dict[str, Any]]: |
| """``warmup_variants`` -> list of variant dicts. ``"default"`` / ``"all"``: :data:`DEFAULT_WARMUP_VARIANTS`; |
| ``None`` / ``"none"`` / ``False`` / ``[]``: none; a dict: one variant; a list of dicts: those.""" |
| if spec is None or spec is False or (isinstance(spec, str) and spec.lower() == "none"): |
| return [] |
| if isinstance(spec, str): |
| if spec.lower() in ("default", "all"): |
| return [dict(v) for v in DEFAULT_WARMUP_VARIANTS] |
| raise ValueError(f"warmup_variants must be 'default', 'none', a dict or a list of dicts, got {spec!r}") |
| if isinstance(spec, dict): |
| return [dict(spec)] |
| out = [] |
| for v in spec: |
| if not isinstance(v, dict): |
| raise ValueError(f"warmup_variants list items must be dicts (model.warmup kwargs), got {v!r}") |
| out.append(dict(v)) |
| return out |
|
|
|
|
| @functools.lru_cache(maxsize=4) |
| def _dummy_frame_cached(W: int, H: int) -> np.ndarray: |
| a = _dummy_frame(W, H, seed=W * 7 + H) |
| a.setflags(write=False) |
| return a |
|
|
|
|
| def _dummy_frame(W: int, H: int, seed: int) -> np.ndarray: |
| """A synthetic ``(H, W, 3)`` uint8 frame with photo-like statistics (smooth shading + texture + noise), so the |
| decoders and the resize run at the cost of a real frame (a flat image decodes much faster).""" |
| g = np.random.default_rng(seed) |
| x = np.arange(W, dtype=np.float32)[None, :, None] |
| y = np.arange(H, dtype=np.float32)[:, None, None] |
| c = np.array([0.0, 1.1, 2.3], np.float32) |
| a = 128 + 90 * np.sin(x / (W / 3.1) + c) * np.cos(y / (H / 2.3) - c) + 25 * np.sin(x / 3.7) * np.cos(y / 5.3) |
| a += g.integers(-24, 25, (H, W, 1), dtype=np.int8) |
| return np.clip(a, 0, 255).astype(np.uint8) |
|
|
|
|
| def _dummy_boxes(W: int, H: int, n: int) -> List[List[float]]: |
| """n hand-sized boxes; the second one crosses the frame edge (exercises the grey padding).""" |
| boxes = [] |
| for i in range(n): |
| if i % 2 == 0: |
| boxes.append([0.30 * W + i, 0.35 * H, 0.55 * W + i, 0.70 * H]) |
| else: |
| boxes.append([-0.05 * W, 0.55 * H, 0.20 * W, 0.95 * H + i]) |
| return boxes |
|
|
|
|
| |
| |
| |
| @dataclass |
| class HandPrediction: |
| """One hand. Field names follow the HTTP ``/predict`` response; :attr:`mano_params` gives the upstream |
| HaMeR ``pred_mano_params`` layout. |
| |
| Attributes: |
| rotmats: ``(16, 3, 3)`` float32 rotation matrices; index 0 = global orientation, 1..15 = MANO hand pose. |
| Right-hand (model) frame: for a left hand they describe the mirrored crop, as upstream HaMeR. |
| betas: ``(10,)`` float32 MANO shape coefficients. |
| cam: ``(3,)`` float32 weak-perspective camera ``[s, tx, ty]`` in crop space (as the network emits it). |
| cam_t_full: ``(3,)`` float64 camera translation ``[tx, ty, tz]`` for the full frame (upstream |
| ``cam_crop_to_full``, focal length :attr:`focal_length`); x mirrored back for a left hand. |
| focal_length: pixels, ``5000 / 256 * max(W, H)``. |
| bbox: ``(x1, y1, x2, y2)`` in original pixels (the request box, or the centre square). |
| bbox_source: ``"request"`` or ``"center_square"``. |
| bbox_center: ``(2,)`` crop centre in original pixels. |
| bbox_size: side of the square crop in original pixels (``rescale_factor * max(w, h)``). |
| is_right, rescale_factor: the request values. |
| image_size: ``(W, H)`` of the frame. |
| vertices: ``(778, 3)`` float64 MANO mesh in metres (MANO frame, x mirrored back for a left hand), or |
| ``None`` without ``MANO_RIGHT.pkl`` or with ``return_vertices=False``. |
| joints: ``(21, 3)`` float64 joints in metres (same frame), or ``None`` without MANO. |
| keypoints_2d: ``(21, 2)`` float64 joints projected into the original image, pixels, or ``None`` without MANO. |
| timing_ms: ``{"preprocess", "model", "mano", "total"}`` wall times of this call. |
| """ |
|
|
| rotmats: np.ndarray |
| betas: np.ndarray |
| cam: np.ndarray |
| cam_t_full: np.ndarray |
| focal_length: float |
| bbox: Tuple[float, float, float, float] |
| bbox_source: str |
| bbox_center: np.ndarray |
| bbox_size: float |
| is_right: bool |
| rescale_factor: float |
| image_size: Tuple[int, int] |
| vertices: Optional[np.ndarray] = None |
| joints: Optional[np.ndarray] = None |
| keypoints_2d: Optional[np.ndarray] = None |
| timing_ms: Dict[str, float] = field(default_factory=dict) |
|
|
| @property |
| def global_orient(self) -> np.ndarray: |
| """``(1, 3, 3)`` global orientation (= ``rotmats[:1]``).""" |
| return self.rotmats[:1] |
|
|
| @property |
| def hand_pose(self) -> np.ndarray: |
| """``(15, 3, 3)`` MANO hand pose (= ``rotmats[1:]``).""" |
| return self.rotmats[1:] |
|
|
| @property |
| def mano_params(self) -> Dict[str, np.ndarray]: |
| """Upstream layout: ``{"global_orient": (1,3,3), "hand_pose": (15,3,3), "betas": (10,)}``.""" |
| return {"global_orient": self.global_orient, "hand_pose": self.hand_pose, "betas": self.betas} |
|
|
| @property |
| def mano_available(self) -> bool: |
| return self.joints is not None |
|
|
| def to_dict(self, rounded: bool = True) -> Dict[str, Any]: |
| """JSON-ready dict with the keys and rounding of the HTTP ``/predict`` response (``timing_ms`` and |
| ``trace_active`` aside). ``rounded=False`` keeps full precision.""" |
| def r(a: Any, nd: int = 6) -> Any: |
| a = np.asarray(a, dtype=np.float64) |
| return (np.round(a, nd) if rounded else a).tolist() |
|
|
| out: Dict[str, Any] = { |
| "image_size": list(self.image_size), |
| "bbox": r(self.bbox, 3), |
| "bbox_source": self.bbox_source, |
| "is_right": self.is_right, |
| "rescale_factor": self.rescale_factor, |
| "bbox_center": r(self.bbox_center, 3), |
| "bbox_size": round(float(self.bbox_size), 3) if rounded else float(self.bbox_size), |
| "focal_length": round(float(self.focal_length), 3) if rounded else float(self.focal_length), |
| "rotmats": r(self.rotmats), |
| "betas": r(self.betas), |
| "cam": r(self.cam), |
| "cam_t_full": r(self.cam_t_full), |
| "mano_available": self.mano_available, |
| } |
| if self.joints is not None: |
| out["joints"] = r(self.joints) |
| out["keypoints_2d"] = r(self.keypoints_2d, 3) |
| if self.vertices is not None: |
| out["vertices"] = r(self.vertices) |
| return out |
|
|
|
|
| |
| |
| |
| def apply_served_config(override: bool = False) -> Dict[str, Tuple[str, str]]: |
| """Put the published ``tt-model.yaml`` serve configuration (:data:`hamer.served_config.SERVED_ENV`) into |
| ``os.environ``. The knobs are read when the device graph is built, so call this before building. With |
| ``override=False`` (default) a value already in the environment is kept; returns ``{name: (kept, served)}`` |
| for every kept value that differs from the published one.""" |
| differ = {} |
| for k, v in SERVED_ENV.items(): |
| cur = os.environ.get(k) |
| if cur is None or override: |
| os.environ[k] = v |
| elif cur != v: |
| differ[k] = (cur, v) |
| return differ |
|
|
|
|
| def resolve_weights(weights_dir: Optional[Union[str, os.PathLike]] = None, repo_id: str = WEIGHTS_REPO, |
| revision: Optional[str] = WEIGHTS_REVISION) -> Tuple[Path, Path]: |
| """``(hamer.ckpt, mano_mean_params.npz)``: from ``weights_dir`` if given, else from the Hugging Face cache |
| (downloaded on first use, 2.7 GB). ``HAMER_WEIGHTS_DIR`` is honoured when ``weights_dir`` is ``None``.""" |
| weights_dir = weights_dir or os.environ.get("HAMER_WEIGHTS_DIR") or None |
| if weights_dir: |
| d = Path(weights_dir).expanduser() |
| ck, mean = d / CKPT_FILE, d / MEAN_FILE |
| for f in (ck, mean): |
| if not f.is_file(): |
| raise FileNotFoundError(f"{d} does not contain {f.name}") |
| return ck, mean |
| from huggingface_hub import hf_hub_download |
|
|
| ck = Path(hf_hub_download(repo_id=repo_id, filename=CKPT_FILE, revision=revision)) |
| mean = Path(hf_hub_download(repo_id=repo_id, filename=MEAN_FILE, revision=revision)) |
| return ck, mean |
|
|
|
|
| def resolve_mano_path(mano_path: Optional[Union[str, os.PathLike]] = None) -> Optional[Path]: |
| """``MANO_RIGHT.pkl`` (licence-bound, not redistributed): ``mano_path`` (file or directory), else |
| ``HAMER_MANO_PATH``, else ``$TT_DIT_CACHE_DIR/MANO_RIGHT.pkl`` or ``~/.cache/tt-model/hamer-p150/weights/``. |
| An explicit path that does not exist raises; the default locations return ``None`` when absent.""" |
| explicit = mano_path or os.environ.get("HAMER_MANO_PATH") |
| if explicit: |
| p = Path(explicit).expanduser() |
| p = p / MANO_FILE if p.is_dir() else p |
| if not p.is_file(): |
| raise FileNotFoundError(f"MANO model not found at {p}") |
| return p |
| cache = os.environ.get("TT_DIT_CACHE_DIR") |
| base = Path(cache) if cache else Path.home() / ".cache" / "tt-model" / "hamer-p150" / "weights" |
| p = base / MANO_FILE |
| return p if p.is_file() else None |
|
|
|
|
| def load_reference(weights_dir: Optional[Union[str, os.PathLike]] = None, repo_id: str = WEIGHTS_REPO, |
| revision: Optional[str] = WEIGHTS_REVISION) -> torch.nn.Module: |
| """The torch CPU reference (``hamer.reference.Hamer``) with the released checkpoint loaded.""" |
| from .checkpoint import load_hamer_checkpoint |
| from .reference import Hamer |
|
|
| ck, mean = resolve_weights(weights_dir, repo_id, revision) |
| ref = Hamer().eval() |
| report = load_hamer_checkpoint(ref, ck, mean) |
| if report["loaded"] == 0 or report["missing_in_ref"] != 0 or report["unexpected_in_ref"] != 0: |
| raise RuntimeError(f"checkpoint {ck} does not map cleanly onto the reference: {report}") |
| if not bool(ref.head.init_cam.abs().sum() > 0): |
| raise RuntimeError(f"{mean} did not populate the MANO mean parameters") |
| return ref |
|
|
|
|
| def open_device(device_id: int = 0, dispatch: str = "auto", verbose: bool = False) -> Tuple[Any, str]: |
| """Open one Blackhole chip the way the published numbers were measured: a 1x1 ``MeshDevice`` with a 200 MiB |
| trace region, ETH dispatch (the full 12x10 Tensix grid computes) and the larger kernel-config ring. |
| |
| ``dispatch``: ``"auto"`` (ETH, falling back to the default worker dispatch with a warning when the tt-metal build |
| cannot place dispatch on Ethernet cores), ``"eth"`` or ``"worker"``. When ``TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE`` |
| is not set and ``ttnn`` is not imported yet, it is set to ``11,9`` (grid capped at 12x10, as on the measured |
| hosts; tt-metal reads it at ``import ttnn``). |
| |
| Returns ``(device, dispatch_used)``. Close it with ``ttnn.close_mesh_device`` (or pass it to |
| :meth:`HamerModel.from_pretrained` as ``device=`` and close it yourself afterwards).""" |
| if dispatch not in ("auto", "eth", "worker"): |
| raise ValueError(f"dispatch must be 'auto', 'eth' or 'worker', got {dispatch!r}") |
| order = ["eth", "worker"] if dispatch == "auto" else [dispatch] |
| |
| |
| if "eth" in order and GRID_OVERRIDE_ENV not in os.environ and "ttnn" not in sys.modules: |
| os.environ[GRID_OVERRIDE_ENV] = f"{SERVED_GRID[0] - 1},{SERVED_GRID[1] - 1}" |
| import ttnn |
|
|
| extra_kb = int(os.environ.get("HAMER_KCFG_EXTRA_KB", "128")) |
| last: Optional[BaseException] = None |
| for i, d in enumerate(order): |
| kw: Dict[str, Any] = {} |
| if d == "eth": |
| kw["dispatch_core_config"] = ttnn.DispatchCoreConfig(ttnn.DispatchCoreType.ETH) |
| if extra_kb > 0: |
| kw["worker_l1_size"] = 1461376 - extra_kb * 1024 |
| try: |
| dev = ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(1, 1), physical_device_ids=[device_id], |
| trace_region_size=TRACE_REGION_BYTES, num_command_queues=1, **kw) |
| except Exception as e: |
| last = e |
| if i + 1 < len(order): |
| warnings.warn( |
| f"hamer: ETH dispatch is not available on this tt-metal build / chip ({e}); falling back to the " |
| "default worker dispatch. The published numbers were measured with ETH dispatch on the 12x10 " |
| "grid; worker dispatch leaves a smaller grid on a p150 and runs slower.", RuntimeWarning, stacklevel=2) |
| continue |
| g = dev.compute_with_storage_grid_size() |
| if (g.x, g.y) != SERVED_GRID: |
| warnings.warn( |
| f"hamer: compute grid is {g.x}x{g.y}; the device graph is tuned and validated on " |
| f"{SERVED_GRID[0]}x{SERVED_GRID[1]} (to cap a larger grid at 12x10, set {GRID_OVERRIDE_ENV}=11,9 before ttnn is " |
| "imported).", |
| RuntimeWarning, stacklevel=2) |
| _log(verbose, f"device {device_id} open: dispatch={d} grid={g.x}x{g.y}") |
| return dev, d |
| raise RuntimeError(f"hamer: could not open device {device_id}: {last}") from last |
|
|
|
|
| def _close_device(dev: Any) -> None: |
| import ttnn |
|
|
| if isinstance(dev, getattr(ttnn, "MeshDevice", ())): |
| ttnn.close_mesh_device(dev) |
| else: |
| ttnn.close_device(dev) |
|
|
|
|
| def _pcc(a: torch.Tensor, b: torch.Tensor) -> float: |
| a = a.reshape(-1).double() - a.double().mean() |
| b = b.reshape(-1).double() - b.double().mean() |
| return float((a @ b) / (a.norm() * b.norm()).clamp_min(1e-12)) |
|
|
|
|
| |
| |
| |
| def to_pil(image: ImageLike) -> Image.Image: |
| """A frame as an RGB ``PIL.Image``. Accepts a file path, a ``PIL.Image`` (any mode), an ``(H, W, 3)`` uint8 |
| RGB array, or a torch tensor ``(3, H, W)`` / ``(H, W, 3)`` (uint8, or float in ``[0, 1]``).""" |
| if isinstance(image, Image.Image): |
| return image if image.mode == "RGB" else image.convert("RGB") |
| if isinstance(image, (str, os.PathLike)): |
| im = Image.open(image) |
| im.load() |
| return im if im.mode == "RGB" else im.convert("RGB") |
| if isinstance(image, torch.Tensor): |
| t = image.detach().cpu() |
| if t.dim() == 4 and t.shape[0] == 1: |
| t = t[0] |
| if t.dim() == 3 and t.shape[0] in (1, 3) and t.shape[-1] not in (1, 3): |
| t = t.permute(1, 2, 0) |
| image = t.numpy() |
| if isinstance(image, np.ndarray): |
| a = image |
| if a.ndim == 2: |
| a = np.stack([a] * 3, -1) |
| if a.ndim != 3 or a.shape[-1] not in (1, 3, 4): |
| raise ValueError(f"expected an (H, W, 3) RGB image array, got shape {a.shape}") |
| if a.shape[-1] == 1: |
| a = np.repeat(a, 3, -1) |
| a = a[..., :3] |
| if a.dtype != np.uint8: |
| if np.issubdtype(a.dtype, np.floating): |
| if a.min() < 0.0 or a.max() > 1.0: |
| raise ValueError("float image arrays must be RGB in [0, 1]; pass uint8 for 0..255 pixels") |
| a = np.round(a * 255.0) |
| a = a.astype(np.uint8) |
| return Image.fromarray(np.ascontiguousarray(a), "RGB") |
| raise TypeError(f"unsupported image type {type(image).__name__}") |
|
|
|
|
| def _check_bbox(bbox: Optional[BBox], img: Image.Image, rescale_factor: float) -> Tuple[Tuple[float, ...], str]: |
| """Same rules as the server's ``_validate_bbox`` (ValueError instead of HTTP 400).""" |
| if bbox is None: |
| return center_square_bbox(img), "center_square" |
| try: |
| x1, y1, x2, y2 = (float(v) for v in bbox) |
| except (TypeError, ValueError): |
| raise ValueError(f"bbox must be four numbers [x1, y1, x2, y2], got {bbox!r}") from None |
| if not all(np.isfinite([x1, y1, x2, y2])) or x2 <= x1 or y2 <= y1: |
| raise ValueError(f"bbox {list(bbox)} is not a valid [x1, y1, x2, y2] with x2 > x1, y2 > y1") |
| if x2 < 0 or y2 < 0 or x1 > img.width or y1 > img.height: |
| raise ValueError(f"bbox {list(bbox)} lies entirely outside the {img.width}x{img.height} image") |
| side = hamer_bbox_size((x1, y1, x2, y2), rescale_factor) |
| if side < MIN_CROP_PIXELS: |
| raise ValueError(f"bbox {list(bbox)} is degenerate: crop side {side:.3f} px is below {MIN_CROP_PIXELS:g} px") |
| return (x1, y1, x2, y2), "request" |
|
|
|
|
| def _is_box_list(bbox: Any) -> bool: |
| """True for several boxes (``[[x1, y1, x2, y2], ...]`` or an ``(N, 4)`` array).""" |
| if bbox is None: |
| return False |
| a = np.asarray(bbox, dtype=object) |
| return a.ndim == 2 or (a.ndim == 1 and len(a) > 0 and isinstance(a[0], (list, tuple, np.ndarray))) |
|
|
|
|
| |
| |
| |
| class HamerModel: |
| """HaMeR (ViT-H/16 + MANO head) on one Blackhole chip. Build with :meth:`from_pretrained`. |
| |
| Call it with a frame and a hand box:: |
| |
| hand = model(image, bbox=[x1, y1, x2, y2], is_right=True) # -> HandPrediction |
| hands = model(image, bbox=[box_a, box_b], is_right=[True, False]) # -> list[HandPrediction] |
| |
| ``image``: file path, ``PIL.Image``, ``(H, W, 3)`` uint8 RGB numpy array, or torch tensor (see :func:`to_pil`). |
| One hand per device call (the device graph is batch 1); several boxes on one frame decode the frame once. |
| Thread-safe (one device call at a time). Release the chip with :meth:`close` or a ``with`` block. |
| """ |
|
|
| def __init__(self, runner: Callable[[torch.Tensor], torch.Tensor], *, ref: torch.nn.Module, |
| raw_input: bool, mano: Optional[torch.nn.Module] = None, tt: Any = None, |
| owned_device: Any = None, info: Optional[Dict[str, Any]] = None) -> None: |
| """Low-level constructor; use :meth:`from_pretrained` or :meth:`from_reference`.""" |
| self._runner = runner |
| self.ref = ref |
| self.raw_input = raw_input |
| self.mano = mano |
| self.tt = tt |
| self._owned_device = owned_device |
| self.info: Dict[str, Any] = dict(info or {}) |
| self._lock = threading.Lock() |
| self._closed = False |
| self._warm: Dict[Tuple[Tuple[str, Any], ...], Dict[str, Any]] = {} |
|
|
| |
| @classmethod |
| def from_pretrained( |
| cls, |
| device_id: int = 0, |
| *, |
| device: Any = None, |
| weights_dir: Optional[Union[str, os.PathLike]] = None, |
| repo_id: str = WEIGHTS_REPO, |
| revision: Optional[str] = WEIGHTS_REVISION, |
| mano_path: Optional[Union[str, os.PathLike]] = None, |
| dispatch: str = "auto", |
| use_trace: bool = True, |
| check_accuracy: bool = True, |
| min_pcc: float = 0.99, |
| served_config: bool = True, |
| warmup_variants: Any = "default", |
| verbose: bool = False, |
| ) -> "HamerModel": |
| """Load the released checkpoint and build the model on a chip, then warm up every call variant in |
| ``warmup_variants`` so that your first real call runs at the steady speed. |
| |
| Args: |
| device_id: chip to open (``TT_VISIBLE_DEVICES`` / ``TT_METAL_VISIBLE_DEVICES`` order). |
| device: an already-opened ttnn device to use instead (not closed by :meth:`close`). Open it with |
| :func:`open_device` to get the published configuration; any 1x1 ``MeshDevice`` with a trace region of |
| at least 200 MiB works. |
| weights_dir: directory with ``hamer.ckpt`` + ``mano_mean_params.npz``; default: the Hugging Face cache |
| (``repo_id`` at ``revision``, downloaded on first use). |
| mano_path: ``MANO_RIGHT.pkl`` (file or directory) for vertices / joints / 2-D keypoints; default: see |
| :func:`resolve_mano_path`. Without it the outputs are the MANO parameters and cameras. |
| dispatch: ``"auto"`` (ETH, else worker with a warning), ``"eth"`` or ``"worker"``. |
| use_trace: capture and replay one ttnn trace (the published path). ``False`` runs eagerly (slow). |
| check_accuracy: run the server's warm-up gate: eager, trace replay and a fresh input vs the torch CPU |
| reference, PCC >= ``min_pcc`` (about 1 s of CPU work). |
| served_config: apply the ``tt-model.yaml`` serve configuration (:func:`apply_served_config`) first. |
| Values already in the environment win and are reported in a warning. |
| warmup_variants: host call variants to prepare after the device warm-up: ``"default"`` (all of |
| :data:`DEFAULT_WARMUP_VARIANTS`, about 0.25 s), ``None`` / ``"none"``, or a list of |
| :meth:`warmup` kwargs dicts. The device graph (one shape) is always compiled and traced. |
| verbose: print the boot steps and keep the tt-metal / ttnn log output. ``False`` (default): errors only |
| (``TT_LOGGER_LEVEL=error``, ``LOGURU_LEVEL=WARNING`` when not set and ``ttnn`` is not imported yet). |
| """ |
| t_boot = time.perf_counter() |
| _quiet_logs(verbose) |
| if served_config: |
| differ = apply_served_config() |
| if differ: |
| warnings.warn("hamer: environment overrides the published configuration: " |
| + ", ".join(f"{k}={a!r} (published {b!r})" for k, (a, b) in differ.items()), |
| RuntimeWarning, stacklevel=2) |
| torch.set_grad_enabled(False) |
| t = time.perf_counter() |
| ref = load_reference(weights_dir, repo_id, revision) |
| _log(verbose, f"checkpoint loaded in {time.perf_counter() - t:.1f}s") |
| mano = _load_mano(mano_path, verbose) |
| info = {"weights": {"repo": None if weights_dir else repo_id, "revision": None if weights_dir else revision, |
| "dir": str(weights_dir) if weights_dir else None}} |
| return cls.from_reference(ref, device_id=device_id, device=device, mano=mano, dispatch=dispatch, |
| use_trace=use_trace, check_accuracy=check_accuracy, min_pcc=min_pcc, |
| warmup_variants=warmup_variants, verbose=verbose, _info=info, _t_boot=t_boot) |
|
|
| @classmethod |
| def from_reference( |
| cls, |
| ref: torch.nn.Module, |
| *, |
| backend: str = "tt", |
| device_id: int = 0, |
| device: Any = None, |
| mano: Optional[torch.nn.Module] = None, |
| dispatch: str = "auto", |
| use_trace: bool = True, |
| check_accuracy: bool = True, |
| min_pcc: float = 0.99, |
| warmup_variants: Any = "auto", |
| verbose: bool = False, |
| _info: Optional[Dict[str, Any]] = None, |
| _t_boot: Optional[float] = None, |
| ) -> "HamerModel": |
| """Build from a torch reference module you already hold (``hamer.reference.Hamer`` with weights loaded). |
| ``backend="cpu"`` runs that module on the host instead of the chip (no ttnn needed; slow; same API). |
| ``warmup_variants``: as in :meth:`from_pretrained`; ``"auto"`` = ``"default"`` on the chip, none on the CPU.""" |
| t_boot = time.perf_counter() if _t_boot is None else _t_boot |
| info = dict(_info or {}) |
| ref.eval() |
| if isinstance(warmup_variants, str) and warmup_variants == "auto": |
| warmup_variants = "default" if backend == "tt" else None |
| variants = _warmup_spec(warmup_variants) |
| for v in variants: |
| normalize_warmup_variant(**v) |
| if backend == "cpu": |
| model = cls(lambda x: ref(x), ref=ref, raw_input=False, mano=mano, info=info) |
| model._warmup_host(variants, verbose) |
| model.info["boot_seconds"] = round(time.perf_counter() - t_boot, 1) |
| return model |
| if backend != "tt": |
| raise ValueError(f"backend must be 'tt' or 'cpu', got {backend!r}") |
|
|
| owned = None |
| if device is None: |
| device, used = open_device(device_id, dispatch, verbose) |
| owned = device |
| else: |
| used = "caller" |
| from .model import build_from_reference |
| tt = None |
| try: |
| t = time.perf_counter() |
| tt = build_from_reference(ref, device_id=device_id, strict=True, use_trace=use_trace, device=device) |
| if tt._vit_params is None or tt._head_params is None: |
| raise RuntimeError("tt-nn parameters were not uploaded") |
| _log(verbose, f"weights uploaded in {time.perf_counter() - t:.1f}s") |
| raw = getattr(tt.fused_opts, "input", "norm") == "u8" |
| g = device.compute_with_storage_grid_size() |
| info.update(backend="tt", device_id=device_id if used != "caller" else None, dispatch=used, |
| num_command_queues=1 if used != "caller" else None, |
| grid=f"{g.x}x{g.y}", fused=tt.fused_opts.describe()) |
| model = cls(tt, ref=ref, raw_input=raw, mano=mano, tt=tt, owned_device=owned, info=info) |
| model.info["warmup"] = model._warmup(check_accuracy, min_pcc, verbose) |
| except BaseException: |
| if tt is not None: |
| tt.close() |
| if owned is not None: |
| _close_device(owned) |
| raise |
| if use_trace and not tt.trace_active: |
| model.close() |
| raise RuntimeError("ttnn trace was not captured during warm-up") |
| try: |
| from . import native |
|
|
| native.load() |
| model.info["native_preprocess"] = native.STATUS |
| except Exception: |
| pass |
| try: |
| model._warmup_host(variants, verbose) |
| except BaseException: |
| model.close() |
| raise |
| model.info["boot_seconds"] = round(time.perf_counter() - t_boot, 1) |
| _log(verbose, f"ready in {model.info['boot_seconds']}s ({model.info['grid']}, dispatch={used})") |
| return model |
|
|
| def _warmup(self, check: bool, min_pcc: float, verbose: bool) -> Dict[str, Any]: |
| """The server's warm-up (``hamer.server.app.lifespan``): forward 1 eager (JIT), forward 2 trace capture + |
| replay on the same input, forward 3 a different input through the copy-into-trace-buffer path.""" |
| from .reference import sample_input |
|
|
| if self.raw_input: |
| g = torch.Generator().manual_seed(1) |
| u_a = torch.randint(0, 256, (1, 3, IMAGE_HEIGHT, IMAGE_WIDTH), generator=g, dtype=torch.uint8) |
| u_b = torch.randint(0, 256, (1, 3, IMAGE_HEIGHT, IMAGE_WIDTH), generator=g, dtype=torch.uint8) |
| x_a, x_b, d_a, d_b = normalize_raw(u_a), normalize_raw(u_b), u_a, u_b |
| else: |
| x_a, x_b = sample_input(batch=1, seed=0), sample_input(batch=1, seed=1) |
| d_a, d_b = x_a, x_b |
| outs = [] |
| with torch.inference_mode(): |
| for i, d in enumerate((d_a, d_a, d_b)): |
| t = time.perf_counter() |
| out = self._runner(d) |
| if tuple(out.shape) != (1, 157) or not torch.isfinite(out).all(): |
| raise RuntimeError(f"warm-up forward {i + 1} produced an invalid output {tuple(out.shape)}") |
| outs.append(out.detach().clone()) |
| _log(verbose, f"warm-up forward {i + 1}/3: {(time.perf_counter() - t) * 1e3:.0f} ms") |
| if not check: |
| return {"checked": False} |
| ref_a, ref_b = self.ref(x_a), self.ref(x_b) |
| res = {"eager_vs_ref_pcc": round(_pcc(outs[0], ref_a), 5), |
| "replay_vs_eager_pcc": round(_pcc(outs[1], outs[0]), 5), |
| "fresh_input_vs_ref_pcc": round(_pcc(outs[2], ref_b), 5), "min_pcc": min_pcc, "checked": True} |
| _log(verbose, f"warm-up PCC {res}") |
| for k in ("eager_vs_ref_pcc", "replay_vs_eager_pcc", "fresh_input_vs_ref_pcc"): |
| if res[k] < min_pcc: |
| raise RuntimeError(f"warm-up gate failed: {k} = {res[k]} < {min_pcc}") |
| return res |
|
|
| def _warmup_host(self, variants: List[Dict[str, Any]], verbose: bool) -> None: |
| """Run the requested host variants (after the device warm-up and the native-helper load).""" |
| self.info.setdefault("warmup_variants", []) |
| if not variants: |
| return |
| t = time.perf_counter() |
| |
| |
| gc.collect() |
| for v in variants: |
| rec = self.warmup(**v) |
| _log(verbose, f"warm-up {rec['variant']}: {rec['ms']:.1f} ms") |
| self.info["warmup_host_seconds"] = round(time.perf_counter() - t, 3) |
|
|
| def warmup(self, call: str = "frame", **variant: Any) -> Dict[str, Any]: |
| """Prepare one call variant: run it once on a dummy input of the real size, so that its one-time host costs |
| (lazy imports, image-decoder init, first native preprocess call on a frame, first MANO kernels, first |
| allocations) are paid now and not in your first real call. Idempotent: a variant that is already warm |
| returns its record without running again. The outputs of the dummy call are checked and dropped (nothing is |
| cached; real calls are not changed). |
| |
| Args (the defaults are the quickstart call): |
| call: ``"frame"`` (``model(image, bbox)``), ``"crops"`` (:meth:`forward_crops`), ``"run_crop"``. |
| input (frame): ``"path"`` (an image file), ``"pil"``, ``"array"`` (uint8 ``(H, W, 3)``), ``"tensor"`` |
| (uint8 ``(3, H, W)``). |
| fmt (frame, path): ``"jpeg"`` or ``"png"`` (the decoder that the first file of that format loads). |
| hands (frame): boxes per call; 1 = single-box form, >= 2 = list form. |
| is_right (frame): ``True``, ``False`` (mirrored crop) or ``"mixed"`` (with ``hands >= 2``). |
| bbox (frame): ``"box"`` or ``"center"`` (``bbox=None``). |
| frame_size (frame): ``(W, H)`` of the dummy frame, default ``(640, 480)``. |
| crop_dtype (crops, run_crop): ``"uint8"`` or ``"float"`` (ImageNet-normalized). |
| crop_width (crops): 256 (upstream crop) or 192. |
| batch (crops): crops per call (one device call each), default 2. |
| |
| Returns ``{"variant": {...}, "ms": wall time, "warmed": True/False (False: it was already warm)}``. |
| """ |
| self._check_open() |
| key = normalize_warmup_variant(call, **variant) |
| done = self._warm.get(key) |
| if done is not None: |
| return dict(done, warmed=False) |
| v = dict(key) |
| t = time.perf_counter() |
| self._run_warmup_variant(v) |
| rec = {"variant": {k: (list(x) if isinstance(x, tuple) else x) for k, x in v.items()}, |
| "ms": round((time.perf_counter() - t) * 1e3, 2)} |
| self._warm[key] = rec |
| self.info.setdefault("warmup_variants", []).append(rec) |
| return dict(rec, warmed=True) |
|
|
| def _run_warmup_variant(self, v: Dict[str, Any]) -> None: |
| call = v["call"] |
| if call == "frame": |
| W, H = v["frame_size"] |
| a = _dummy_frame_cached(W, H).copy() |
| n = v["hands"] |
| boxes = _dummy_boxes(W, H, n) if v["bbox"] == "box" else None |
| if n == 1: |
| bbox: Any = None if boxes is None else boxes[0] |
| right: Any = bool(v["is_right"]) |
| else: |
| bbox = boxes if boxes is not None else [None] * n |
| right = [bool(i % 2 == 0) for i in range(n)] if v["is_right"] == "mixed" else [bool(v["is_right"])] * n |
| inp = v["input"] |
| with tempfile.TemporaryDirectory(prefix="hamer-warmup-") as d: |
| if inp == "path": |
| image: Any = Path(d) / ("frame.jpg" if v["fmt"] == "jpeg" else "frame.png") |
| Image.fromarray(a).save(image, **({"quality": 90} if v["fmt"] == "jpeg" else {"compress_level": 1})) |
| elif inp == "pil": |
| image = Image.fromarray(a) |
| elif inp == "array": |
| image = a |
| else: |
| image = torch.from_numpy(a).permute(2, 0, 1) |
| if n > 1 and boxes is None: |
| img = to_pil(image) |
| bbox = [list(center_square_bbox(img))] * n |
| out = self(image, bbox, is_right=right) |
| hands = out if isinstance(out, list) else [out] |
| if len(hands) != n: |
| raise RuntimeError(f"warm-up {v}: {len(hands)} outputs for {n} boxes") |
| for h in hands: |
| if not (np.isfinite(h.rotmats).all() and np.isfinite(h.cam_t_full).all()): |
| raise RuntimeError(f"warm-up {v}: non-finite output") |
| h.to_dict() |
| |
| |
| |
| side = max(W, H) / 2.0 |
| preprocess_pil(Image.fromarray(a), (W / 2 - side / 2, H / 2 - side / 2, W / 2 + side / 2, H / 2 + side / 2), |
| rescale_factor=2.0, is_right=True, raw=self.raw_input, contiguous=not self.raw_input) |
| return |
| width = v.get("crop_width", IMAGE_WIDTH) |
| B = v.get("batch", 1) |
| u = torch.from_numpy(_dummy_frame(width, IMAGE_HEIGHT * B, seed=width + B).copy()) |
| u = u.reshape(B, IMAGE_HEIGHT, width, 3).permute(0, 3, 1, 2).contiguous() |
| x = u if v["crop_dtype"] == "uint8" else normalize_raw(u) |
| if call == "crops": |
| res = self.forward_crops(x) |
| if not torch.isfinite(res["pred_cam"]).all(): |
| raise RuntimeError(f"warm-up {v}: non-finite output") |
| return |
| out = self.run_crop(x) |
| if tuple(out.shape) != (1, 157) or not torch.isfinite(out).all(): |
| raise RuntimeError(f"warm-up {v}: invalid output {tuple(out.shape)}") |
|
|
| @property |
| def warmed_variants(self) -> List[Dict[str, Any]]: |
| """The variants prepared so far (by :meth:`from_pretrained` and :meth:`warmup`).""" |
| return [dict(k) for k in self._warm] |
|
|
| |
| @property |
| def device(self) -> Any: |
| """The ttnn device (``None`` for the CPU backend or after :meth:`close`).""" |
| return None if self.tt is None else self.tt.device |
|
|
| @property |
| def trace_active(self) -> bool: |
| return bool(getattr(self.tt, "trace_active", False)) |
|
|
| @property |
| def mano_available(self) -> bool: |
| return self.mano is not None |
|
|
| @property |
| def faces(self) -> Optional[np.ndarray]: |
| """``(1538, 3)`` int64 MANO triangle indices (``None`` without ``MANO_RIGHT.pkl``).""" |
| return None if self.mano is None else np.asarray(self.mano.faces, dtype=np.int64) |
|
|
| |
| def __call__( |
| self, |
| image: ImageLike, |
| bbox: Optional[Union[BBox, Sequence[BBox], np.ndarray]] = None, |
| *, |
| is_right: Union[bool, Sequence[bool]] = True, |
| rescale_factor: float = 2.0, |
| return_vertices: bool = True, |
| ) -> Union[HandPrediction, List[HandPrediction]]: |
| """Hand mesh recovery for one hand (``bbox`` = ``[x1, y1, x2, y2]``) or several hands of the same frame |
| (``bbox`` = a list / ``(N, 4)`` array of boxes, ``is_right`` = one bool or one per box). |
| |
| Args: |
| image: the FULL frame (path, PIL, uint8 ``(H, W, 3)`` RGB array, or torch tensor; see :func:`to_pil`). |
| bbox: hand box(es) in original pixels. ``None``: the centre square of the frame (the camera is then |
| geometrically meaningless; for smoke tests). |
| is_right: ``False`` for a left hand: the crop is mirrored, the outputs are mirrored back. |
| rescale_factor: square crop side = ``rescale_factor * max(box w, h)``, in ``(0.5, 5.0]``; default 2.0. |
| return_vertices: include the 778 mesh vertices when MANO is available. |
| |
| Returns: |
| :class:`HandPrediction` for one box, ``list[HandPrediction]`` for a list of boxes. |
| """ |
| t0 = time.perf_counter() |
| img = to_pil(image) |
| if not (0.5 < float(rescale_factor) <= 5.0): |
| raise ValueError(f"rescale_factor must be in (0.5, 5.0], got {rescale_factor}") |
| if _is_box_list(bbox): |
| boxes = [list(b) for b in np.asarray(bbox, dtype=np.float64).reshape(-1, 4)] |
| rights = list(is_right) if isinstance(is_right, (list, tuple, np.ndarray)) else [is_right] * len(boxes) |
| if len(rights) != len(boxes): |
| raise ValueError(f"{len(boxes)} boxes but {len(rights)} is_right values") |
| return [self._predict_one(img, b, bool(r), float(rescale_factor), return_vertices, t0 if i == 0 else None) |
| for i, (b, r) in enumerate(zip(boxes, rights))] |
| if isinstance(is_right, (list, tuple, np.ndarray)): |
| raise ValueError("is_right must be a single bool for a single box") |
| return self._predict_one(img, bbox, bool(is_right), float(rescale_factor), return_vertices, t0) |
|
|
| predict = __call__ |
|
|
| def _predict_one(self, img: Image.Image, bbox: Optional[BBox], is_right: bool, rescale_factor: float, |
| return_vertices: bool, t0: Optional[float]) -> HandPrediction: |
| """The server's ``/predict`` body after the image decode (same calls, same order, no rounding).""" |
| t0 = time.perf_counter() if t0 is None else t0 |
| box, source = _check_bbox(bbox, img, rescale_factor) |
| W, H = img.size |
| raw_in = self.raw_input |
| x, (cx, cy, bbox_size) = preprocess_pil(img, box, rescale_factor=rescale_factor, is_right=is_right, |
| raw=raw_in, contiguous=not raw_in) |
| if tuple(x.shape) != (1, 3, IMAGE_HEIGHT, IMAGE_WIDTH): |
| raise RuntimeError(f"preprocess produced {tuple(x.shape)}") |
| t1 = time.perf_counter() |
| out = self.run_crop(x if raw_in else x.contiguous()) |
| t2 = time.perf_counter() |
|
|
| out = out[0].float().cpu() |
| rotmats = out[:144].reshape(16, 3, 3) |
| betas = out[144:154] |
| cam = out[154:157] |
| focal = FOCAL_LENGTH / MODEL_IMAGE_SIZE * max(W, H) |
| cam_full = cam.numpy().astype(np.float64).copy() |
| if not is_right: |
| cam_full[1] *= -1.0 |
| cam_t = cam_crop_to_full(cam_full, np.array([cx, cy]), bbox_size, W, H, focal) |
| pred = HandPrediction( |
| rotmats=rotmats.numpy().copy(), betas=betas.numpy().copy(), cam=cam.numpy().copy(), cam_t_full=cam_t, |
| focal_length=float(focal), bbox=tuple(float(v) for v in box), bbox_source=source, |
| bbox_center=np.array([cx, cy], dtype=np.float64), bbox_size=float(bbox_size), is_right=is_right, |
| rescale_factor=rescale_factor, image_size=(W, H), |
| ) |
| t3 = t2 |
| if self.mano is not None: |
| with torch.inference_mode(): |
| verts, joints = (v[0] for v in self.mano.forward_with_joints(rotmats[None], betas[None])) |
| kp2d, _ = project_mesh(joints, cam, (cx, cy), bbox_size, W, H, is_right=is_right) |
| verts_np = verts.numpy().astype(np.float64).copy() |
| joints_np = joints.numpy().astype(np.float64).copy() |
| if not is_right: |
| verts_np[:, 0] *= -1.0 |
| joints_np[:, 0] *= -1.0 |
| pred.joints, pred.keypoints_2d = joints_np, np.asarray(kp2d, dtype=np.float64) |
| if return_vertices: |
| pred.vertices = verts_np |
| t3 = time.perf_counter() |
| pred.timing_ms = {"preprocess": (t1 - t0) * 1e3, "model": (t2 - t1) * 1e3, "mano": (t3 - t2) * 1e3, |
| "total": (time.perf_counter() - t0) * 1e3} |
| return pred |
|
|
| def run_crop(self, crop: torch.Tensor) -> torch.Tensor: |
| """One device call on one network-input crop -> ``(1, 157)`` float32 |
| ``[rotmats (16*9) | betas (10) | cam (3)]``. |
| |
| ``crop``: ``(1, 3, 256, 192)``, either uint8 pixels (``hamer.preprocess.preprocess_pil(..., raw=True)``) or |
| the ImageNet-normalized float crop (``raw=False``). With the published uint8 device input a float crop must |
| be an exactly normalized 8-bit image (else ``ValueError``).""" |
| self._check_open() |
| if self.tt is None and crop.dtype == torch.uint8: |
| crop = normalize_raw(crop) |
| with self._lock, torch.inference_mode(): |
| return self._runner(crop) |
|
|
| def forward_crops(self, img: torch.Tensor) -> Dict[str, Any]: |
| """Upstream-style batch forward on pre-cropped hands (``hamer.models.HAMER.forward(batch)`` conventions). |
| |
| Args: |
| img: ``(B, 3, 256, 256)`` (the upstream ``ViTDetDataset`` crop; the model uses columns 32:224) or |
| ``(B, 3, 256, 192)``; uint8 pixels or ImageNet-normalized float (``(x / 255 - mean) / std``). |
| Left hands must already be mirrored, as upstream. |
| |
| Returns (torch tensors, batch first; B device calls, one per crop): |
| ``pred_cam`` (B, 3), ``pred_cam_t`` (B, 3) (crop camera, focal 5000 px for the 256 crop), |
| ``focal_length`` (B, 2), ``pred_mano_params`` = ``{"global_orient": (B, 1, 3, 3), |
| "hand_pose": (B, 15, 3, 3), "betas": (B, 10)}``, and with MANO ``pred_vertices`` (B, 778, 3) and |
| ``pred_keypoints_3d`` (B, 21, 3) in metres. |
| """ |
| if img.dim() != 4 or img.shape[1] != 3 or img.shape[2] != IMAGE_HEIGHT or img.shape[3] not in (256, IMAGE_WIDTH): |
| raise ValueError(f"expected (B, 3, 256, 256) or (B, 3, 256, 192), got {tuple(img.shape)}") |
| if img.shape[3] == 256: |
| img = img[:, :, :, W_SLICE:-W_SLICE] |
| outs = [self.run_crop(img[b:b + 1].contiguous())[0].float().clone() for b in range(img.shape[0])] |
| out = torch.stack(outs) |
| B = out.shape[0] |
| rot = out[:, :144].reshape(B, 16, 3, 3) |
| betas, cam = out[:, 144:154], out[:, 154:157] |
| cam_t = torch.stack([cam[:, 1], cam[:, 2], 2 * FOCAL_LENGTH / (MODEL_IMAGE_SIZE * cam[:, 0] + 1e-9)], dim=-1) |
| res: Dict[str, Any] = { |
| "pred_cam": cam, "pred_cam_t": cam_t, |
| "focal_length": torch.full((B, 2), FOCAL_LENGTH, dtype=torch.float32), |
| "pred_mano_params": {"global_orient": rot[:, :1], "hand_pose": rot[:, 1:], "betas": betas}, |
| } |
| if self.mano is not None: |
| with torch.inference_mode(): |
| verts, joints = self.mano.forward_with_joints(rot, betas) |
| res["pred_vertices"], res["pred_keypoints_3d"] = verts, joints |
| return res |
|
|
| |
| def _check_open(self) -> None: |
| if self._closed: |
| raise RuntimeError("HamerModel is closed") |
|
|
| def close(self) -> None: |
| """Release the trace and the device tensors, and close the chip if this object opened it. Idempotent.""" |
| if self._closed: |
| return |
| self._closed = True |
| with self._lock: |
| if self.tt is not None: |
| self.tt.close() |
| if self._owned_device is not None: |
| try: |
| _close_device(self._owned_device) |
| finally: |
| self._owned_device = None |
|
|
| def __enter__(self) -> "HamerModel": |
| return self |
|
|
| def __exit__(self, *exc: Any) -> None: |
| self.close() |
|
|
| def __del__(self) -> None: |
| try: |
| self.close() |
| except Exception: |
| pass |
|
|
| def __repr__(self) -> str: |
| state = "closed" if self._closed else f"backend={self.info.get('backend')}" |
| extra = "" if self.tt is None else f", grid={self.info.get('grid')}, dispatch={self.info.get('dispatch')}" |
| return f"HamerModel({state}{extra}, mano={'yes' if self.mano is not None else 'no'})" |
|
|
|
|
| def _load_mano(mano_path: Optional[Union[str, os.PathLike]], verbose: bool) -> Optional[torch.nn.Module]: |
| p = resolve_mano_path(mano_path) |
| if p is None: |
| _log(verbose, "MANO_RIGHT.pkl not found: MANO parameters and cameras only (pass mano_path= for the mesh)") |
| return None |
| from .mano import MANO |
|
|
| m = MANO(p) |
| m.eval() |
| _log(verbose, f"MANO loaded from {p}") |
| return m |
|
|
