Download code/hamer/preprocess.py from changh95/hamer-p150: direct link, hf CLI and curl.
- Browser
- Download file 11.8 kB
-
https://huggingface.co/changh95/hamer-p150/resolve/main/code/hamer/preprocess.py
- Command line
-
hf download hf://changh95/hamer-p150/code/hamer/preprocess.py
-
curl -L -o preprocess.py https://huggingface.co/changh95/hamer-p150/resolve/main/code/hamer/preprocess.py
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 | |