hamer-p150 / code /hamer /preprocess.py
changh95's picture
Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms
58ba6df verified
Raw History Blame Contribute Delete
11.8 kB
# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Image preprocessing for the HaMeR / HaMeR regressor input.
Mirrors the upstream ``ViTDetDataset`` + ``HAMER.forward`` behaviour:
1. Around the supplied hand bbox, expand by ``rescale_factor`` (default 2.0)
and take the larger side → a square bbox of ``bbox_size`` pixels in
original-image coords centred on the bbox.
2. Crop that square out of the image, resize to ``256 × 256``.
3. Centre-slice ``[..., 32:-32]`` along the W axis → ``256 × 192``. This
matches HaMeR's ``self.backbone(x[:,:,:,32:-32])``.
4. ImageNet-normalize.
"""
from __future__ import annotations
import os
import threading
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Optional, Tuple
import numpy as np
import torch
from PIL import Image
IMAGE_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
IMAGE_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
IMAGE_HEIGHT = 256 # ViT input H after the [:,:,:,32:-32] slice
IMAGE_WIDTH = 192 # ViT input W after the [:,:,:,32:-32] slice
BBOX_CROP_SIZE = 256 # square the bbox is resized to before W-slice
W_SLICE = 32 # pixels removed from each side of W
# round 13: the 192 x 256 resize runs as RESIZE_BANDS horizontal bands of output rows, the first on the calling thread and
# the others on a small thread pool (PIL releases the GIL while resampling). Bit-identical to the one-call resize: the
# crop side S is an integer, so every band's box edge r * S / 256 and PIL's per-row filter centres box0 + (y + 0.5) * scale
# are exact dyadic floats -> the same filter windows and fixed-point weights as the full resize (checked on 1000 random
# crops x 2 band counts, tests/test_fused_host.py). NOT the default: a standalone loop gains 0.3-0.5 ms (2 bands), but
# inside the server (300-request A/B on chip 5) the preprocess p50 moved 2.03 / 2.01 -> 1.87 / 2.33 ms (2 bands) and
# 2.41 ms (3 bands) -- the pool threads wake cold on the loaded host. HAMER_PRE_BANDS=1 (default): one call.
RESIZE_BANDS = max(1, int(os.environ.get("HAMER_PRE_BANDS", "1")))
_POOL: Optional[ThreadPoolExecutor] = None
_POOL_LOCK = threading.Lock()
# round 14 (P1): native crop + resize (hamer/native, bit-identical to PIL, ~3x faster); HAMER_PRE_NATIVE=0 -> PIL path.
# Not used together with HAMER_PRE_BANDS > 1 (that knob measures the PIL band path).
def native_enabled() -> bool:
return RESIZE_BANDS == 1 and os.environ.get("HAMER_PRE_NATIVE", "1").strip().lower() not in ("0", "false", "no", "off")
def _pool(n: int) -> ThreadPoolExecutor:
global _POOL
with _POOL_LOCK:
if _POOL is None or _POOL._max_workers < n: # noqa: SLF001
_POOL = ThreadPoolExecutor(max_workers=n, thread_name_prefix="hamer-resize")
return _POOL
def resize_crop_hw(crop: Image.Image, is_right: bool = True, bands: Optional[int] = None) -> np.ndarray:
"""The square crop -> uint8 HWC ``[256, 192, 3]``: the centre 192 of the 256 resized columns (BILINEAR), mirrored
for a left hand. ``bands`` > 1 splits the output rows over threads (bit-identical, see RESIZE_BANDS)."""
S = crop.width
x0, x1 = W_SLICE * S / BBOX_CROP_SIZE, (BBOX_CROP_SIZE - W_SLICE) * S / BBOX_CROP_SIZE
nb = RESIZE_BANDS if bands is None else max(1, int(bands))
if nb == 1:
c = crop.resize((IMAGE_WIDTH, BBOX_CROP_SIZE), Image.BILINEAR, box=(x0, 0, x1, S))
if not is_right:
c = c.transpose(Image.FLIP_LEFT_RIGHT)
return np.asarray(c).copy()
out = np.empty((BBOX_CROP_SIZE, IMAGE_WIDTH, 3), np.uint8)
edges = [round(i * BBOX_CROP_SIZE / nb) for i in range(nb + 1)]
def band(i: int) -> None:
r0, r1 = edges[i], edges[i + 1]
b = crop.resize((IMAGE_WIDTH, r1 - r0), Image.BILINEAR,
box=(x0, r0 * S / BBOX_CROP_SIZE, x1, r1 * S / BBOX_CROP_SIZE))
if not is_right:
b = b.transpose(Image.FLIP_LEFT_RIGHT)
out[r0:r1] = np.asarray(b)
futs = [_pool(nb - 1).submit(band, i) for i in range(1, nb)]
band(0)
for f in futs:
f.result()
return out
def hamer_bbox_size(bbox_xyxy: Tuple[float, float, float, float], rescale_factor: float = 2.0) -> float:
"""Mirrors HaMeR's ``ViTDetDataset`` scale computation:
scale = rescale_factor * (bbox_w, bbox_h) / 200
bbox_size = (scale * 200).max()
With ``BBOX_SHAPE=None`` (the released checkpoint's default), no
aspect-ratio expansion is applied — the square bbox just fits the
longer side.
"""
x1, y1, x2, y2 = bbox_xyxy
w, h = x2 - x1, y2 - y1
return float(rescale_factor * max(w, h))
def crop_bbox(
image: Image.Image,
bbox_xyxy: Tuple[float, float, float, float],
rescale_factor: float = 2.0,
) -> Tuple[Image.Image, Tuple[float, float, float]]:
"""Crop the HaMeR-style square hand bbox out of ``image`` and return the
256×256 PIL crop plus ``(cx, cy, bbox_size)`` in *original-image coords*.
"""
crop, meta = _square_crop(image, bbox_xyxy, rescale_factor)
return crop.resize((BBOX_CROP_SIZE, BBOX_CROP_SIZE), Image.BILINEAR), meta
def _square_crop(
image: Image.Image,
bbox_xyxy: Tuple[float, float, float, float],
rescale_factor: float = 2.0,
) -> Tuple[Image.Image, Tuple[float, float, float]]:
"""The HaMeR-style square crop (grey padding) at source resolution, plus ``(cx, cy, bbox_size)``."""
x1, y1, x2, y2 = bbox_xyxy
cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0
bbox_size = hamer_bbox_size(bbox_xyxy, rescale_factor)
half = bbox_size / 2.0
crop_box = (cx - half, cy - half, cx + half, cy + half)
# Pad with grey if the bbox extends beyond the image edges.
crop = Image.new("RGB", (int(round(bbox_size)), int(round(bbox_size))), (128, 128, 128))
src = image.crop((max(0, int(round(crop_box[0]))), max(0, int(round(crop_box[1]))),
min(image.width, int(round(crop_box[2]))), min(image.height, int(round(crop_box[3])))))
paste_x = max(0, int(round(-crop_box[0])))
paste_y = max(0, int(round(-crop_box[1])))
crop.paste(src, (paste_x, paste_y))
return crop, (cx, cy, bbox_size)
def preprocess(
image_path: str | Path,
bbox_xyxy: Optional[Tuple[float, float, float, float]] = None,
rescale_factor: float = 2.0,
is_right: bool = True,
raw: bool = False,
) -> Tuple[torch.Tensor, Optional[Tuple[float, float, float]]]:
"""Load image, crop HaMeR-style around ``bbox_xyxy``, return
``(input_tensor, bbox_meta)``.
``bbox_xyxy`` is in original-image pixel coords. When omitted, the
image is centre-cropped to a square — useful for smoke-tests, but
geometrically meaningless for the camera prediction.
``is_right=False`` mirrors the crop horizontally before normalisation —
HaMeR is trained on right hands only, so a left hand is fed in
flipped form (matches ``ViTDetDataset``'s ``flip = (right == 0)``).
``bbox_meta = (cx, cy, bbox_size)`` lets the visualiser reverse the
crop to overlay the mesh back onto the original frame.
"""
img = Image.open(image_path).convert("RGB")
return preprocess_pil(img, bbox_xyxy, rescale_factor=rescale_factor, is_right=is_right, raw=raw)
def center_square_bbox(img: Image.Image) -> Tuple[float, float, float, float]:
"""The centre square of ``img`` as an xyxy bbox (the no-bbox fallback)."""
side = min(img.width, img.height)
return (
(img.width - side) / 2.0,
(img.height - side) / 2.0,
(img.width + side) / 2.0,
(img.height + side) / 2.0,
)
def preprocess_pil(
img: Image.Image,
bbox_xyxy: Optional[Tuple[float, float, float, float]] = None,
rescale_factor: float = 2.0,
is_right: bool = True,
raw: bool = False,
contiguous: bool = True,
) -> Tuple[torch.Tensor, Tuple[float, float, float]]:
"""In-memory variant of :func:`preprocess` for an already-decoded PIL
image (what a server holds). Same crop / flip / slice / normalise
sequence, byte-for-byte. ``raw=True``: the uint8 ``[1, 3, 256, 192]`` crop before the
normalisation (``normalize_raw`` of it is bit-identical to ``raw=False``); the
``TT_FUSED_INPUT=u8`` device graph takes it directly. ``contiguous=False`` (raw only): return the
``[1, 3, 256, 192]`` view of the HWC array instead of a CHW copy (the server; same values).
"""
if img.mode != "RGB":
img = img.convert("RGB")
if bbox_xyxy is None:
bbox_xyxy = center_square_bbox(img)
res = None
if native_enabled():
# round 14 (P1): crop + grey pad + box resize + mirror in one native call, bit-identical to the PIL path below
# (hamer/native/resize.c; self-checked at load, tests/test_fused_host.py::test_native_crop_resize_bit_identical)
from hamer import native
# planar CHW straight from the native pass: the raw crop is contiguous [1, 3, 256, 192] at no extra cost (the
# model's centring then reads it contiguously: 0.02 vs 0.07 ms for the strided HWC view)
# round 15: the frame's pixel buffer is read in place (HAMER_PRE_ZEROCOPY=1; no np.asarray copy)
res = native.crop_resize_image(img, bbox_xyxy, rescale_factor, is_right, chw=True)
if res is not None:
u8_chw, bbox_meta = res
x = torch.from_numpy(u8_chw).unsqueeze(0)
if raw:
return x, bbox_meta
return (x.float() / 255.0 - IMAGE_MEAN) / IMAGE_STD, bbox_meta
crop, bbox_meta = _square_crop(img, bbox_xyxy, rescale_factor)
# Only the centre 192 of the 256 resized columns are used: resize just that box (round 6: bit-identical to
# resize-to-256 + slice -- PIL's filter windows depend only on the output pixel centre and the source image,
# checked on 400 random crops -- and ~25% less host time). The flip commutes with the symmetric slice.
return _finish(torch.from_numpy(resize_crop_hw(crop, is_right)), bbox_meta, raw, contiguous)
def _finish(u8: torch.Tensor, bbox_meta, raw: bool, contiguous: bool):
if raw:
x = u8.permute(2, 0, 1).unsqueeze(0)
# round 13 (server): contiguous=False keeps the view (no CHW copy); the model's input centring reads it strided
return (x.contiguous() if contiguous else x), bbox_meta
arr = u8.float() / 255.0
arr = arr.permute(2, 0, 1).unsqueeze(0) # (1, 3, 256, 192)
return (arr - IMAGE_MEAN) / IMAGE_STD, bbox_meta
def normalize_raw(u8: torch.Tensor) -> torch.Tensor:
"""uint8 ``[B, 3, 256, 192]`` crop (``preprocess_pil(raw=True)``) -> the normalized fp32 input,
bit-identical to ``preprocess_pil(raw=False)``."""
return (u8.float() / 255.0 - IMAGE_MEAN) / IMAGE_STD
def denormalize(x: torch.Tensor) -> np.ndarray:
img = x * IMAGE_STD + IMAGE_MEAN
img = img.squeeze(0).clamp(0, 1).permute(1, 2, 0).numpy()
return (img * 255).astype(np.uint8)
def map_network_to_original(
pts_192x256: np.ndarray,
bbox_meta: Tuple[float, float, float],
) -> np.ndarray:
"""Map 2-D points in network-input pixel coords (192 × 256) back to
original-image pixel coords.
Inverse of the preprocess: shift x by ``W_SLICE`` to undo the slice,
then scale by ``bbox_size / BBOX_CROP_SIZE`` and translate by the bbox
top-left corner.
"""
cx, cy, bbox_size = bbox_meta
pts = pts_192x256.astype(np.float64).copy()
pts[:, 0] += W_SLICE # 192-wide → 256-wide
pts *= bbox_size / BBOX_CROP_SIZE # 256×256 → bbox_size×bbox_size
pts[:, 0] += cx - bbox_size / 2.0
pts[:, 1] += cy - bbox_size / 2.0
return pts