"""HaMeR image features: the MANO decoder token of each hand, one pair per frame. The hand crop (_crop_hand) is adapted from HaMeR's hamer/datasets/utils.py (https://github.com/geopavlakos/hamer, MIT License, Copyright (c) 2023 UC Regents, Georgios Pavlakos), see THIRD_PARTY_NOTICES.md. """ from pathlib import Path import cv2 import numpy as np import torch from hamer.configs import get_config from hamer.models import HAMER from torch.utils.data import DataLoader from tqdm import tqdm def hand_boxes_from_keypoints(kp2d, threshold, min_kpts=4): """Per-frame boxes around the confident keypoints of one hand. Args: kp2d: (L, J, 3) hand keypoints [x, y, score] threshold: keypoints with score > threshold are used min_kpts: frames with fewer confident keypoints get an all-zero box (= no hand) Returns: (L, 4) boxes [x1, y1, x2, y2] """ x, y = kp2d[..., 0], kp2d[..., 1] mask = kp2d[..., 2] > threshold inf = x.new_tensor(float("inf")) boxes = torch.stack( [ torch.where(mask, x, inf).amin(dim=1), torch.where(mask, y, inf).amin(dim=1), torch.where(mask, x, -inf).amax(dim=1), torch.where(mask, y, -inf).amax(dim=1), ], dim=1, ) valid = mask.sum(dim=1) >= min_kpts return torch.where(valid.unsqueeze(1), boxes, torch.zeros_like(boxes)) class HamerHandFeatures: """Per-frame HaMeR tokens (1024-d) of both hands, given their bounding boxes. The token (the MANO decoder output consumed by HaMeR's regression heads) is read with a forward hook. Hands are cropped like in HaMeR's demo; the left hand is mirrored, as HaMeR only models right hands. Frames without a hand box get an all-black crop. Args: data_dir: HaMeR's data directory (from its hamer_demo_data download): contains hamer_ckpts/checkpoints/hamer.ckpt, hamer_ckpts/model_config.yaml, data/mano_mean_params.npz, and the MANO model data/mano/MANO_RIGHT.pkl fp16: run the network under float16 autocast (~1.9x faster, tokens nearly identical) """ TOKEN_DIM = 1024 def __init__(self, data_dir, fp16=True, batch_size=16, num_workers=8, device="cuda"): data_dir = Path(data_dir) cfg = get_config(str(data_dir / "hamer_ckpts" / "model_config.yaml")) cfg.defrost() cfg.MANO.MODEL_PATH = str(data_dir / "data" / "mano") cfg.MANO.MEAN_PARAMS = str(data_dir / "data" / "mano_mean_params.npz") if cfg.MODEL.BACKBONE.TYPE == "vit" and "BBOX_SHAPE" not in cfg.MODEL: cfg.MODEL.BBOX_SHAPE = [192, 256] cfg.MODEL.BACKBONE.pop("PRETRAINED_WEIGHTS", None) cfg.freeze() self.cfg = cfg # like hamer.models.load_hamer, without the renderer; built on the GPU: initializing the # weights on the CPU takes ~2 s, and the checkpoint replaces them anyway with torch.device(device): self.model = HAMER.load_from_checkpoint( str(data_dir / "hamer_ckpts" / "checkpoints" / "hamer.ckpt"), strict=False, cfg=cfg, init_renderer=False, map_location="cpu", # ZeroGPU: deserialize on CPU; .to(device) handles the move ) self.model.eval().to(device) self.fp16, self.batch_size, self.num_workers, self.device = ( fp16, batch_size, num_workers, device ) # fmt: skip self._tokens = None self.model.mano_head.transformer.register_forward_hook(self._store_tokens) def _store_tokens(self, _module, _inputs, output): self._tokens = output.squeeze(1) @torch.no_grad() def __call__(self, frame_dir, hand_bbx_xyxy): """frame_dir with %06d.jpg frames, hand_bbx_xyxy (L, 2, 4) [left, right] boxes -> (L, 2, 1024) tokens""" dataset = _FrameDataset(Path(frame_dir), hand_bbx_xyxy, self.cfg) loader = DataLoader( dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, pin_memory=True, ) tokens = torch.zeros((len(dataset), 2, self.TOKEN_DIM), dtype=torch.float32) i = 0 with tqdm(total=len(dataset), desc="HaMeR hand features", unit="frame") as pbar: for crops in loader: bs = crops.shape[0] tokens[i : i + bs] = ( self._forward(crops.reshape(-1, 3, 256, 256)) .reshape(bs, 2, self.TOKEN_DIM) .cpu() ) i += bs pbar.update(bs) return tokens def _forward(self, crops): """(B, 3, 256, 256) normalized crops -> (B, 1024) tokens""" with torch.autocast(self.device, dtype=torch.float16, enabled=self.fp16): # HaMeR's ViT backbone sees the central 256x192 of the square crop feats = self.model.backbone(crops[:, :, :, 32:-32].to(self.device)) self.model.mano_head(feats) # the hook stores the tokens return self._tokens.float() class _FrameDataset(torch.utils.data.Dataset): """Normalized (2, 3, S, S) [left, right] hand crops per frame.""" def __init__(self, frame_dir, hand_bbx_xyxy, cfg): self.frame_dir, self.hand_bbx_xyxy = frame_dir, hand_bbx_xyxy self.img_size = cfg.MODEL.IMAGE_SIZE self.mean = 255.0 * np.array(cfg.MODEL.IMAGE_MEAN) self.std = 255.0 * np.array(cfg.MODEL.IMAGE_STD) self.bbox_shape = cfg.MODEL.BBOX_SHAPE def __len__(self): return len(self.hand_bbx_xyxy) def __getitem__(self, idx): path = self.frame_dir / f"{idx:06d}.jpg" img = cv2.imread(str(path)) if img is None: raise FileNotFoundError(f"Could not read frame {path}") boxes = self.hand_bbx_xyxy[idx].numpy() center = (boxes[:, 2:4] + boxes[:, 0:2]) / 2.0 scale = 2.0 * (boxes[:, 2:4] - boxes[:, 0:2]) / 200.0 # box enlarged 2x, in 200 px crops = [] for i in range(2): # left, right if scale[i, 0] == 0 or scale[i, 1] == 0: # no hand in this frame crops.append(np.zeros((1, 3, self.img_size, self.img_size), dtype=np.float32)) continue bbox_size = _expand_to_aspect_ratio(scale[i] * 200, self.bbox_shape).max() patch = _crop_hand( img.copy(), center[i, 0], center[i, 1], bbox_size, self.img_size, flip=i == 0 ) patch = patch[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) # BGR HWC -> RGB CHW for c in range(min(img.shape[2], 3)): patch[c] = (patch[c] - self.mean[c]) / self.std[c] crops.append(patch[None]) return np.concatenate(crops, axis=0) def _expand_to_aspect_ratio(size_wh, target_aspect_ratio): """Grow a (w, h) box size to the target (w, h) aspect ratio.""" w, h = size_wh w_t, h_t = target_aspect_ratio if h / w < h_t / w_t: return np.array([w, max(w * h_t / w_t, h)]) return np.array([max(h * w_t / h_t, w), h]) def _crop_hand(img, c_x, c_y, bbox_size, patch_size, flip): """Square crop of size bbox_size around (c_x, c_y), resized to patch_size (optionally mirrored), with black borders.""" if flip: img = img[:, ::-1, :] c_x = img.shape[1] - c_x - 1 src_center = np.array([c_x, c_y], dtype=np.float64) src_down = np.array([0, bbox_size * 0.5], dtype=np.float32) src_right = np.array([bbox_size * 0.5, 0], dtype=np.float32) dst_center = np.array([patch_size * 0.5, patch_size * 0.5], dtype=np.float32) dst_down = np.array([0, patch_size * 0.5], dtype=np.float32) dst_right = np.array([patch_size * 0.5, 0], dtype=np.float32) src = np.zeros((3, 2), dtype=np.float32) src[0, :] = src_center src[1, :] = src_center + src_down src[2, :] = src_center + src_right dst = np.zeros((3, 2), dtype=np.float32) dst[0, :] = dst_center dst[1, :] = dst_center + dst_down dst[2, :] = dst_center + dst_right trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) return cv2.warpAffine( img, trans, (int(patch_size), int(patch_size)), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0, )