""" Online HaMeR pipeline — YOLO detection + full HaMeR inference. Starting from raw images, a single forward pass simultaneously outputs: - backbone features (ViT-H 192 tokens × 1280D) - 2D landmarks (crop coordinate system [0,1]) - crop intrinsics - detection confidence Reuses the crop utility functions and coordinate conversion logic from preprocessing/hamer_crop_and_skeleton.py. """ from __future__ import annotations import os import sys from pathlib import Path from typing import Optional import cv2 import numpy as np import torch # Project root directory _PROJECT_ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(_PROJECT_ROOT)) from utils.dual_conf import select_target_box_dual_conf # HaMeR is a submodule located at third_party/hamer; checkpoint path is set via environment variable (see README) HAMER_ROOT = _PROJECT_ROOT / "third_party" / "hamer" HAMER_CKPT = Path(os.environ.get( "HAMER_CKPT", str(HAMER_ROOT / "_DATA" / "hamer_ckpts" / "checkpoints" / "hamer.ckpt"), )) # _DATA: prefer environment variable; otherwise inferred from HAMER_CKPT (ckpt is at <_DATA>/hamer_ckpts/checkpoints/) HAMER_DATA = Path(os.environ.get( "HAMER_DATA", str(HAMER_CKPT.parents[2]) if os.environ.get("HAMER_CKPT") else str(HAMER_ROOT / "_DATA"), )) DETECTOR_PT = Path(os.environ["DETECTOR_CKPT"]) if "DETECTOR_CKPT" in os.environ else None # WiLoR hand detector class mapping: cls 0=left, 1=right _SIDE_TO_CLS = {"left": 0, "right": 1} # ── crop utility functions (imported from preprocessing) ───────────────────────────────── def _import_crop_utils(): """Lazy import to avoid import errors in non-GPU environments.""" sys.path.insert(0, str(_PROJECT_ROOT)) from preprocessing.hamer_crop_and_skeleton import ( expand_to_aspect_ratio, gen_trans_from_patch_cv, compute_crop_intrinsics, _cam_crop_to_full, _project_full_img, ) return expand_to_aspect_ratio, gen_trans_from_patch_cv, compute_crop_intrinsics, _cam_crop_to_full, _project_full_img class OnlineHaMeRPipeline: """ Online HaMeR pipeline: raw image -> YOLO detection -> crop -> full HaMeR inference. Captures backbone features via a forward hook; a single forward pass jointly produces backbone features + 2D landmarks + crop intrinsics + confidence. """ # ImageNet normalization (x255 variant, consistent with the HAMER model) _IMAGENET_MEAN = torch.tensor([123.675, 116.28, 103.53]) _IMAGENET_STD = torch.tensor([58.395, 57.12, 57.375]) def __init__( self, hamer_ckpt: str | None = None, detector_ckpt: str | None = None, device: str = "cuda", det_conf: float = 0.3, hamer_batch_size: int = 32, image_size: int = 256, ): self.device = device self.det_conf = det_conf self.hamer_batch_size = hamer_batch_size self.image_size = image_size # Lazy-import utility functions ( self._expand_to_aspect_ratio, self._gen_trans_from_patch_cv, self._compute_crop_intrinsics, self._cam_crop_to_full, self._project_full_img, ) = _import_crop_utils() # Load YOLO + HAMER _hamer = hamer_ckpt or str(HAMER_CKPT) _detector = detector_ckpt or (str(DETECTOR_PT) if DETECTOR_PT else None) if _detector is None: raise RuntimeError( "YOLO detector weights not specified: please export DETECTOR_CKPT=/path/to/detector.pt, " "or pass the detector_ckpt argument (weights come from the WiLoR repo)" ) self.detector, self.model, self.model_cfg = self._setup_models( _hamer, _detector, device, ) # backbone hook storage self._backbone_features: Optional[torch.Tensor] = None self._hook_handle = self.model.backbone.register_forward_hook( self._backbone_hook ) # normalization buffer self.register_buffer = lambda name, val: setattr(self, name, val.to(device)) self._mean = self._IMAGENET_MEAN.view(1, 3, 1, 1).to(device) self._std = self._IMAGENET_STD.view(1, 3, 1, 1).to(device) def _backbone_hook(self, module, input, output): """Forward hook: captures backbone output (B, C, Hp, Wp) -> (B, 192, 1280).""" feat_map = output # (B, 1280, Hp, Wp) B, C, Hp, Wp = feat_map.shape self._backbone_features = feat_map.flatten(2).transpose(1, 2) # (B, Hp*Wp, 1280) # ── model initialization ──────────────────────────────────────────────── @staticmethod def _setup_models(hamer_ckpt: str, detector_ckpt: str, device: str): """Load the YOLO detector + the full HAMER model.""" import gc sys.path.insert(0, str(HAMER_ROOT)) os.environ["HAMER_DATA"] = str(HAMER_DATA) # Temporarily patch torch.load for compatibility with legacy ckpts _orig_load = torch.load def _patched_load(*a, **k): k.setdefault("weights_only", False) return _orig_load(*a, **k) torch.load = _patched_load from ultralytics import YOLO detector = YOLO(str(detector_ckpt)) detector.to(device) torch.load = _orig_load # restore import hamer.configs as _hcfg _hcfg.CACHE_DIR_HAMER = str(HAMER_DATA) from hamer.configs import get_config from hamer.models.hamer import HAMER model_cfg_path = str(Path(hamer_ckpt).parent.parent / "model_config.yaml") model_cfg = get_config(model_cfg_path, update_cachedir=True) if (model_cfg.MODEL.BACKBONE.TYPE == "vit") and ("BBOX_SHAPE" not in model_cfg.MODEL): model_cfg.defrost() model_cfg.MODEL.BBOX_SHAPE = [192, 256] model_cfg.freeze() if "PRETRAINED_WEIGHTS" in model_cfg.MODEL.BACKBONE: model_cfg.defrost() model_cfg.MODEL.BACKBONE.pop("PRETRAINED_WEIGHTS") model_cfg.freeze() # Manually load to CPU to avoid PL load_from_checkpoint keeping the entire ckpt on GPU model = HAMER(cfg=model_cfg, init_renderer=False) ckpt_dict = torch.load(str(hamer_ckpt), map_location="cpu", weights_only=False) state_dict = ckpt_dict.get("state_dict", ckpt_dict) # Clean up the ckpt dict, keeping only the state_dict del ckpt_dict gc.collect() model.load_state_dict(state_dict, strict=False) del state_dict gc.collect() model = model.to(device).eval() torch.cuda.empty_cache() # Print actual GPU memory usage if torch.cuda.is_available(): free, total = torch.cuda.mem_get_info(0) print(f"[OnlineHaMeR] YOLO + HAMER loaded on {device}, " f"GPU memory: {(total-free)/1024**2:.0f} / {total/1024**2:.0f} MB") return detector, model, model_cfg # ── main interface ───────────────────────────────────────────────────── @torch.no_grad() def process_sequence( self, raw_images_bgr: list[np.ndarray], raw_intrinsics_4d: list[np.ndarray], target_side: str = "right", ) -> dict: """ Process the raw images of a full sequence. Args: raw_images_bgr: list of (H, W, 3) BGR uint8 images raw_intrinsics_4d: list of (4,) [fx, fy, cx, cy] raw intrinsics target_side: "right"/"left", lock onto a single hand -- only select the YOLO box of this hand side; detections of the other hand are always ignored (prevents frames with both hands co-occurring from sliding onto the wrong hand). WiLoR: 0=left, 1=right. Returns: dict with: crop_images: (T, 3, 256, 256) uint8 backbone_features: (T, 192, 1280) float32 — ViT backbone output hamer_landmarks: (T, 21, 2) float32 — crop coordinate system [0,1] crop_intrinsics: (T, 4) float32 hamer_confidence: (T,) float32 detection_valid: (T,) bool det_time_s: float — YOLO detection time hamer_time_s: float — HaMeR inference time """ import time T = len(raw_images_bgr) dev = self.device img_size = self.image_size # ── 1. YOLO detection ──────────────────────────────────────────── t0 = time.perf_counter() detection_batch_size = 32 yolo_out = [] for s in range(0, T, detection_batch_size): e = min(s + detection_batch_size, T) yolo_out.extend( self.detector(raw_images_bgr[s:e], conf=self.det_conf, verbose=False) ) # Parse detection results valid_fi = [] # frame indices with successful detections valid_bboxes = [] # (1, 4) bbox valid_confs = [] # float valid_hw = [] # (H, W) valid_sides = [] # hand side "left"/"right" for each valid frame detection_valid = np.zeros(T, dtype=bool) for fi, yres in enumerate(yolo_out): if len(yres.boxes) == 0: continue clses = yres.boxes.cls.cpu().numpy() confs = yres.boxes.conf.cpu().numpy() bboxes = yres.boxes.xyxy.cpu().numpy() # Lock onto target_side: among boxes of this hand side only, pick the highest conf (WiLoR: 0=left, 1=right) # Detections of the other hand are always ignored, preventing frames with both hands co-occurring from sliding onto the wrong hand target_cls = _SIDE_TO_CLS[target_side] # Pick highest-conf box of target hand side; but if it heavily overlaps (IoU>thr) # a box of the other side with higher conf, the target box is defeated (actually the # other hand) -> skip this frame (cmask will fallback to no-info at inference). sel = select_target_box_dual_conf(clses, confs, bboxes, target_cls=target_cls) if sel is None: continue # target not detected, or defeated by overlapping other-hand box valid_fi.append(fi) valid_bboxes.append(bboxes[sel:sel + 1].astype(np.float32)) valid_confs.append(float(confs[sel])) valid_hw.append(raw_images_bgr[fi].shape[:2]) valid_sides.append(target_side) # after locking, always the target hand side detection_valid[fi] = True t_det = time.perf_counter() # ── 2. generate crops + affine transform ───────────────────────────── BBOX_SHAPE = self.model_cfg.MODEL.get("BBOX_SHAPE", None) rescale_factor = 2.0 crop_bgr_list = [] crop_intr_list = [] box_center_list = [] bbox_size_list = [] for vi, fi in enumerate(valid_fi): img = raw_images_bgr[fi] bbox = valid_bboxes[vi] box = bbox[0].astype(np.float32) center = (box[2:4] + box[0:2]) / 2.0 scale = rescale_factor * (box[2:4] - box[0:2]) / 200.0 bbox_size = float( self._expand_to_aspect_ratio(scale * 200, target_aspect_ratio=BBOX_SHAPE).max() ) center_x, center_y = center trans = self._gen_trans_from_patch_cv( center_x, center_y, bbox_size, bbox_size, img_size, img_size, 1.0, 0, ) crop_bgr = cv2.warpAffine( img, trans, (img_size, img_size), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, ) crop_bgr_list.append(crop_bgr) box_center_list.append(center.astype(np.float32)) bbox_size_list.append(float(bbox_size)) # crop intrinsics intr = raw_intrinsics_4d[fi] if fi < len(raw_intrinsics_4d) else np.zeros(4, np.float32) crop_intr_list.append( self._compute_crop_intrinsics(trans, intr) ) # ── 3. HaMeR batched inference ─────────────────────────────────────── backbone_features_all = np.zeros((T, 192, 1280), dtype=np.float32) landmarks_all = np.zeros((T, 21, 2), dtype=np.float32) crop_intrinsics_all = np.zeros((T, 4), dtype=np.float32) confidence_all = np.zeros(T, dtype=np.float32) crop_images_all = np.zeros((T, 3, 256, 256), dtype=np.uint8) for vi, fi in enumerate(valid_fi): crop_intrinsics_all[fi] = crop_intr_list[vi] confidence_all[fi] = valid_confs[vi] # crop image: BGR → RGB → (3, H, W) crop_images_all[fi] = crop_bgr_list[vi][:, :, ::-1].transpose(2, 0, 1) if len(valid_fi) > 0: box_center_t = torch.tensor( np.stack(box_center_list), dtype=torch.float32, device=dev, ) bbox_size_t = torch.tensor( bbox_size_list, dtype=torch.float32, device=dev, ) img_size_t = torch.tensor( [[float(w), float(h)] for h, w in valid_hw], dtype=torch.float32, device=dev, ) scaled_fl = float( self.model_cfg.EXTRA.FOCAL_LENGTH / self.model_cfg.MODEL.IMAGE_SIZE * img_size_t[0].max().item() ) # Batched inference for s in range(0, len(crop_bgr_list), self.hamer_batch_size): e = min(s + self.hamer_batch_size, len(crop_bgr_list)) crops_np = np.stack(crop_bgr_list[s:e])[:, :, :, ::-1].astype(np.float32) crops_t = torch.from_numpy(crops_np.copy()).permute(0, 3, 1, 2).to(dev) crops_t = (crops_t - self._mean) / self._std out = self.model({"img": crops_t}) # backbone features (captured by the hook, output for the current sub-batch) bf = self._backbone_features.cpu().numpy() # (batch, 192, 1280) for bi in range(e - s): backbone_features_all[valid_fi[s + bi]] = bf[bi] # 2D landmarks: 3D joints -> full image -> normalize -> transform to crop pred_cam = out["pred_cam"] joints_3d = out["pred_keypoints_3d"].detach().cpu().numpy() cam_t_full = self._cam_crop_to_full( pred_cam, box_center_t[s:e], bbox_size_t[s:e], img_size_t[s:e], scaled_fl, ).detach().cpu().numpy() for bi in range(e - s): fi = valid_fi[s + bi] H, W = valid_hw[s + bi] # Project onto the original image kp2d = self._project_full_img( joints_3d[bi], cam_t_full[bi], scaled_fl, W, H, ) # Normalize to [0,1] (original image coordinates) kp_n_orig = np.clip( kp2d / np.array([W, H], dtype=np.float32), 0.0, 1.0, ).astype(np.float32) # Transform to crop coordinate system [0,1] K_crop = crop_intrinsics_all[fi] K_orig = raw_intrinsics_4d[fi] if fi < len(raw_intrinsics_4d) else np.zeros(4) kp_n_crop = self._transform_landmarks_to_crop( kp_n_orig, K_orig, K_crop, W, H, ) landmarks_all[fi] = kp_n_crop # Save box_center/box_size/img_size/sides as full-length-T arrays (invalid frames are 0/default) box_center_all = np.zeros((T, 2), dtype=np.float32) bbox_size_all = np.zeros(T, dtype=np.float32) img_size_all = np.zeros((T, 2), dtype=np.float32) bbox_xyxy_all = np.zeros((T, 4), dtype=np.float32) conf_all = np.zeros(T, dtype=np.float32) sides_all = ["right"] * T for vi, fi in enumerate(valid_fi): box_center_all[fi] = box_center_list[vi] bbox_size_all[fi] = bbox_size_list[vi] img_size_all[fi] = [float(valid_hw[vi][1]), float(valid_hw[vi][0])] # (W, H) bbox_xyxy_all[fi] = valid_bboxes[vi][0] conf_all[fi] = valid_confs[vi] sides_all[fi] = valid_sides[vi] t_hamer = time.perf_counter() return { "crop_images": torch.from_numpy(crop_images_all), "sides": sides_all, "bbox_xyxy": torch.from_numpy(bbox_xyxy_all), "bbox_conf": torch.from_numpy(conf_all), "backbone_features": torch.from_numpy(backbone_features_all), "hamer_landmarks": torch.from_numpy(landmarks_all), "crop_intrinsics": torch.from_numpy(crop_intrinsics_all), "hamer_confidence": torch.from_numpy(confidence_all), "detection_valid": torch.from_numpy(detection_valid), "box_center": torch.from_numpy(box_center_all), "box_size": torch.from_numpy(bbox_size_all), "img_size": torch.from_numpy(img_size_all), "det_time_s": t_det - t0, "hamer_time_s": t_hamer - t_det, } # ── coordinate conversion (consistent with _transform_landmarks_to_crop in dataset.py) ── @staticmethod def _transform_landmarks_to_crop( landmarks_orig_n: np.ndarray, K_orig: np.ndarray, K_crop: np.ndarray, W: int, H: int, ) -> np.ndarray: """Original-image normalized [0,1] -> crop normalized [0,1].""" fx, fy, cx, cy = K_orig fxc, fyc, cxc, cyc = K_crop sx = fxc / fx tx = cxc - fxc * cx / fx sy = fyc / fy ty = cyc - fyc * cy / fy pts_x = landmarks_orig_n[:, 0] * W pts_y = landmarks_orig_n[:, 1] * H crop_x = sx * pts_x + tx crop_y = sy * pts_y + ty result = np.stack([crop_x / 256.0, crop_y / 256.0], axis=-1) return result.astype(np.float32)