| from __future__ import annotations |
|
|
| import os |
| from typing import Optional, Sequence, Tuple |
|
|
| import torch |
| import torch.distributed as dist |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch import Tensor |
| import warnings |
| import logging |
| import numpy as np |
| import cv2 |
| from safetensors.torch import load_file |
| from huggingface_hub import hf_hub_download |
|
|
| class _EfficientNetBackbone(nn.Module): |
| _LAST_CHANNELS: dict[str, int] = { |
| 'b0': 1280, 'b1': 1280, 'b2': 1408, |
| 'b3': 1536, 'b4': 1792, 'b5': 2048, |
| 'b6': 2304, 'b7': 2560, |
| } |
|
|
| def __init__( |
| self, |
| variant: str = 'b5', |
| pretrained: bool = False, |
| out_indices: Tuple[int, ...] = (8,), |
| frozen_stages: int = -1, |
| norm_eval: bool = False, |
| ) -> None: |
| super().__init__() |
|
|
| variant = variant.lower() |
| assert variant in self._LAST_CHANNELS, ( |
| f"Unknown EfficientNet variant '{variant}'. " |
| f"Choose from {list(self._LAST_CHANNELS)}." |
| ) |
|
|
| self.variant = variant |
| self.out_indices = out_indices |
| self.frozen_stages = frozen_stages |
| self.norm_eval = norm_eval |
| self.out_channels = self._LAST_CHANNELS[variant] |
| import torchvision.models as tvm |
|
|
| weights_arg = 'DEFAULT' if pretrained else None |
| builder = getattr(tvm, f'efficientnet_{variant}') |
| is_dist = dist.is_available() and dist.is_initialized() |
| local_rank = int(os.environ.get('LOCAL_RANK', 0)) |
| if is_dist and local_rank != 0: |
| dist.barrier() |
| tv_model = builder(weights=weights_arg) |
| if is_dist and local_rank == 0: |
| dist.barrier() |
|
|
| self.features: nn.Sequential = tv_model.features |
| self.classifier = tv_model.classifier |
|
|
| self._freeze_stages() |
|
|
| def _freeze_stages(self) -> None: |
| for i, layer in enumerate(self.features): |
| if i <= self.frozen_stages: |
| layer.eval() |
| for param in layer.parameters(): |
| param.requires_grad = False |
|
|
| def train(self, mode: bool = True) -> 'EfficientNetBackbone': |
| super().train(mode) |
| self._freeze_stages() |
| if mode and self.norm_eval: |
| for m in self.modules(): |
| if isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)): |
| m.eval() |
| return self |
|
|
| def forward(self, x: Tensor) -> Tuple[Tensor, ...]: |
| outs = [] |
| for i, layer in enumerate(self.features): |
| x = layer(x) |
| if i in self.out_indices: |
| outs.append(x) |
| return tuple(outs) |
|
|
| class HeatmapHead(nn.Module): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| deconv_out_channels: Sequence[int] = (256, 256, 256), |
| deconv_kernel_sizes: Sequence[int] = (4, 4, 4), |
| conv_out_channels: Optional[Sequence[int]] = None, |
| conv_kernel_sizes: Optional[Sequence[int]] = None, |
| final_kernel_size: int = 1, |
| ) -> None: |
| super().__init__() |
|
|
| self.in_channels = in_channels |
| self.out_channels = out_channels |
|
|
| if deconv_out_channels: |
| assert len(deconv_out_channels) == len(deconv_kernel_sizes), ( |
| "'deconv_out_channels' and 'deconv_kernel_sizes' must have " |
| "equal length." |
| ) |
| self.deconv_layers = self._make_deconv_layers( |
| in_channels, deconv_out_channels, deconv_kernel_sizes |
| ) |
| in_channels = deconv_out_channels[-1] |
| else: |
| self.deconv_layers = nn.Identity() |
|
|
| if conv_out_channels: |
| assert conv_kernel_sizes is not None and len( |
| conv_out_channels) == len(conv_kernel_sizes), ( |
| "'conv_out_channels' and 'conv_kernel_sizes' must have " |
| "equal length." |
| ) |
| self.conv_layers = self._make_conv_layers( |
| in_channels, conv_out_channels, conv_kernel_sizes |
| ) |
| in_channels = conv_out_channels[-1] |
| else: |
| self.conv_layers = nn.Identity() |
|
|
| pad = (final_kernel_size - 1) // 2 |
| self.final_layer = nn.Conv2d( |
| in_channels, out_channels, |
| kernel_size=final_kernel_size, |
| padding=pad, |
| ) |
|
|
| self._init_weights() |
|
|
| def _init_weights(self) -> None: |
| for m in self.modules(): |
| if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): |
| nn.init.normal_(m.weight, std=0.001) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
| elif isinstance(m, nn.BatchNorm2d): |
| nn.init.ones_(m.weight) |
| nn.init.zeros_(m.bias) |
|
|
| @staticmethod |
| def _make_deconv_layers( |
| in_channels: int, |
| out_channels_list: Sequence[int], |
| kernel_sizes: Sequence[int], |
| ) -> nn.Sequential: |
| layers: list[nn.Module] = [] |
| for out_ch, ks in zip(out_channels_list, kernel_sizes): |
| if ks == 4: |
| padding, output_padding = 1, 0 |
| elif ks == 3: |
| padding, output_padding = 1, 1 |
| elif ks == 2: |
| padding, output_padding = 0, 0 |
| else: |
| raise ValueError( |
| f"Unsupported deconv kernel size {ks}. Use 2, 3, or 4." |
| ) |
| layers += [ |
| nn.ConvTranspose2d( |
| in_channels, out_ch, |
| kernel_size=ks, stride=2, |
| padding=padding, output_padding=output_padding, |
| bias=False, |
| ), |
| nn.BatchNorm2d(out_ch), |
| nn.ReLU(inplace=True), |
| ] |
| in_channels = out_ch |
| return nn.Sequential(*layers) |
|
|
| @staticmethod |
| def _make_conv_layers( |
| in_channels: int, |
| out_channels_list: Sequence[int], |
| kernel_sizes: Sequence[int], |
| ) -> nn.Sequential: |
| layers: list[nn.Module] = [] |
| for out_ch, ks in zip(out_channels_list, kernel_sizes): |
| padding = (ks - 1) // 2 |
| layers += [ |
| nn.Conv2d(in_channels, out_ch, |
| kernel_size=ks, stride=1, padding=padding), |
| nn.BatchNorm2d(out_ch), |
| nn.ReLU(inplace=True), |
| ] |
| in_channels = out_ch |
| return nn.Sequential(*layers) |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| x = self.deconv_layers(x) |
| x = self.conv_layers(x) |
| x = self.final_layer(x) |
| return x |
|
|
| class EfficientNetB5PoseNet(nn.Module): |
| def __init__( |
| self, |
| num_keypoints: int = 17, |
| pretrained: bool = False, |
| frozen_stages: int = -1, |
| norm_eval: bool = False, |
| deconv_out_channels: Tuple[int, ...] = (256, 256, 256), |
| deconv_kernel_sizes: Tuple[int, ...] = (4, 4, 4), |
| ) -> None: |
| super().__init__() |
|
|
| self.backbone = _EfficientNetBackbone( |
| variant='b5', |
| pretrained=pretrained, |
| out_indices=(8,), |
| frozen_stages=frozen_stages, |
| norm_eval=norm_eval, |
| ) |
| backbone_out_ch = self.backbone.out_channels |
| self.head = HeatmapHead( |
| in_channels=backbone_out_ch, |
| out_channels=num_keypoints, |
| deconv_out_channels=deconv_out_channels, |
| deconv_kernel_sizes=deconv_kernel_sizes, |
| ) |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| feats: Tuple[Tensor, ...] = self.backbone(x) |
| feat: Tensor = feats[-1] |
| heatmaps: Tensor = self.head(feat) |
| return heatmaps |
|
|
|
|
| DEFAULT_INPUT_SIZE = (192, 256) |
|
|
| class PoseEstimator: |
| def __init__(self, model_name, num_keypoints=17, device=None, input_size=DEFAULT_INPUT_SIZE): |
| if device is None: |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.device = torch.device(device) |
| self.input_size = input_size |
| self.model_name = model_name |
| self.model = EfficientNetB5PoseNet(num_keypoints=num_keypoints) |
| if os.path.isfile(model_name): |
| weights_path = model_name |
| elif os.path.isdir(model_name): |
| weights_path = os.path.join(model_name, "model.safetensors") |
| else: |
| weights_path = hf_hub_download(repo_id=model_name, filename="model.safetensors") |
| state_dict = load_file(weights_path, device=str(self.device)) |
| self.model.load_state_dict(state_dict, strict=False) |
| self.model.to(self.device) |
| self.model.eval() |
| self.num_keypoints = num_keypoints |
|
|
| @staticmethod |
| def _get_centers_and_scales_xyxy(person_boxes, scale_factor=1.0): |
| centers, scales = [], [] |
| for box in person_boxes: |
| x1, y1, x2, y2 = box |
| x1, x2 = sorted([x1, x2]) |
| y1, y2 = sorted([y1, y2]) |
| centers.append([(x1+x2)/2.0, (y1+y2)/2.0]) |
| w, h = x2-x1, y2-y1 |
| scales.append([(w/200.0)*scale_factor, (h/200.0)*scale_factor]) |
| return np.array(centers), np.array(scales) |
|
|
| @staticmethod |
| def _process_image(image, bbox, target_size, angle=0, flip=False): |
| try: |
| if image is None or not isinstance(image, np.ndarray): |
| raise ValueError("Invalid image input.") |
| x1, y1, x2, y2 = map(lambda v: int(round(v)), bbox) |
| if x2-x1 <= 0 or y2-y1 <= 0: |
| raise ValueError(f"Invalid bbox: {{bbox}}") |
| x1, y1 = max(0, x1), max(0, y1) |
| x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2) |
| if x2 <= x1 or y2 <= y1: |
| raise ValueError("Invalid bbox after clamping.") |
| cropped = image[y1:y2, x1:x2] |
| resized = cv2.resize(cropped, target_size) |
| if angle != 0: |
| center = (target_size[0]//2, target_size[1]//2) |
| rot = cv2.getRotationMatrix2D(center, angle, 1.0) |
| resized = cv2.warpAffine(resized, rot, target_size) |
| if flip: |
| resized = cv2.flip(resized, 1) |
| return resized, True |
| except Exception: |
| blank = np.zeros((target_size[1], target_size[0], 3), dtype=np.uint8) |
| return blank, False |
|
|
| @staticmethod |
| def _process(image, target_size=(192, 256), angle=0, flip=False, conf_threshold=0.5, model_weights="yolov8n.pt"): |
| try: |
| from ultralytics import YOLO |
| except ImportError: |
| raise ImportError("ultralytics is required. pip install ultralytics") |
| model = YOLO(model_weights) |
| crops, metadata = [], [] |
| if image is None or not isinstance(image, np.ndarray): |
| raise ValueError("Invalid image input.") |
| results = model(image, conf=conf_threshold, classes=[0], verbose=False) |
| bboxes = [] |
| for r in results: |
| for box in r.boxes: |
| bboxes.append(box.xyxy[0].cpu().numpy().tolist()) |
| for idx, bbox in enumerate(bboxes): |
| processed, success = PoseEstimator._process_image(image, bbox, target_size, angle, flip) |
| crops.append(processed) |
| metadata.append({"bbox": bbox, "person_index": idx, "success": success}) |
| if not crops: |
| return None, metadata |
| batch = np.stack(crops, axis=0).transpose(0, 3, 1, 2) |
| return np.ascontiguousarray(batch), metadata |
|
|
| def _preprocess(self, image_bgr): |
| batch, meta = self._process(image_bgr, target_size=self.input_size) |
| if batch is None: |
| return None, meta |
| t = torch.tensor(batch, dtype=torch.float32) / 255.0 |
| return t.to(self.device), meta |
|
|
| @staticmethod |
| def _taylor(heatmap, coord): |
| H, W = heatmap.shape[:2] |
| px, py = int(coord[0]), int(coord[1]) |
| if 1 < px < W-2 and 1 < py < H-2: |
| dx = 0.5*(heatmap[py][px+1]-heatmap[py][px-1]) |
| dy = 0.5*(heatmap[py+1][px]-heatmap[py-1][px]) |
| dxx = 0.25*(heatmap[py][px+2]-2*heatmap[py][px]+heatmap[py][px-2]) |
| dxy = 0.25*(heatmap[py+1][px+1]-heatmap[py-1][px+1]-heatmap[py+1][px-1]+heatmap[py-1][px-1]) |
| dyy = 0.25*(heatmap[py+2][px]-2*heatmap[py][px]+heatmap[py-2][px]) |
| derivative = np.array([[dx],[dy]]) |
| hessian = np.array([[dxx,dxy],[dxy,dyy]]) |
| if dxx*dyy - dxy**2 != 0: |
| offset = -np.linalg.inv(hessian) @ derivative |
| coord += np.squeeze(offset.T, axis=0) |
| return coord |
|
|
| @staticmethod |
| def _get_max_preds(heatmaps): |
| N, K, _, W = heatmaps.shape |
| reshaped = heatmaps.reshape((N, K, -1)) |
| idx = np.argmax(reshaped, 2).reshape((N, K, 1)) |
| maxvals = np.amax(reshaped, 2).reshape((N, K, 1)) |
| preds = np.tile(idx, (1, 1, 2)).astype(np.float32) |
| preds[:,:,0] = preds[:,:,0] % W |
| preds[:,:,1] = np.floor(preds[:,:,1] / W) |
| preds = np.where(np.tile(maxvals, (1, 1, 2)) > 0.0, preds, -1) |
| return preds, maxvals |
|
|
| @staticmethod |
| def _gaussian_blur(heatmaps, kernel=11): |
| border = (kernel-1)//2 |
| B, J, H, W = heatmaps.shape |
| for i in range(B): |
| for j in range(J): |
| origin_max = np.max(heatmaps[i,j]) |
| dr = np.zeros((H+2*border, W+2*border), dtype=np.float32) |
| dr[border:-border, border:-border] = heatmaps[i,j].copy() |
| dr = cv2.GaussianBlur(dr, (kernel, kernel), 0) |
| heatmaps[i,j] = dr[border:-border, border:-border].copy() |
| heatmaps[i,j] *= origin_max / np.max(heatmaps[i,j]) |
| return heatmaps |
|
|
| @staticmethod |
| def transform_preds(coords, center, scale, output_size, use_udp=False): |
| scale = scale * 200.0 |
| if use_udp: |
| sx = scale[0]/(output_size[0]-1.0) |
| sy = scale[1]/(output_size[1]-1.0) |
| else: |
| sx = scale[0]/output_size[0] |
| sy = scale[1]/output_size[1] |
| tc = np.ones_like(coords) |
| tc[:,0] = coords[:,0]*sx + center[0] - scale[0]*0.5 |
| tc[:,1] = coords[:,1]*sy + center[1] - scale[1]*0.5 |
| return tc |
|
|
| @staticmethod |
| def keypoints_from_heatmaps(heatmaps, center, scale, unbiased=False, post_process="default", kernel=11, use_udp=False, target_type="GaussianHeatmap"): |
| heatmaps = heatmaps.copy() |
| if unbiased: |
| assert post_process not in [False, None, "megvii"] |
| if post_process == "default" and unbiased: |
| post_process = "unbiased" |
| if post_process == "megvii": |
| heatmaps = PoseEstimator._gaussian_blur(heatmaps, kernel=kernel) |
| N, K, H, W = heatmaps.shape |
| preds, maxvals = PoseEstimator._get_max_preds(heatmaps) |
| if post_process == "unbiased": |
| heatmaps = np.log(np.maximum(PoseEstimator._gaussian_blur(heatmaps, kernel), 1e-10)) |
| for n in range(N): |
| for k in range(K): |
| preds[n][k] = PoseEstimator._taylor(heatmaps[n][k], preds[n][k]) |
| elif post_process is not None and post_process != "megvii": |
| for n in range(N): |
| for k in range(K): |
| hm = heatmaps[n][k] |
| px, py = int(preds[n][k][0]), int(preds[n][k][1]) |
| if 1 < px < W-1 and 1 < py < H-1: |
| diff = np.array([hm[py][px+1]-hm[py][px-1], hm[py+1][px]-hm[py-1][px]]) |
| preds[n][k] += np.sign(diff)*0.25 |
| for i in range(N): |
| preds[i] = PoseEstimator.transform_preds(preds[i], center[i], scale[i], [W, H], use_udp=use_udp) |
| if post_process == "megvii": |
| maxvals = maxvals/255.0 + 0.5 |
| return preds, maxvals |
|
|
| @torch.no_grad() |
| def predict(self, image_bgr): |
| tensor, meta = self._preprocess(image_bgr) |
| if tensor is None: |
| return np.array([]), np.array([]) |
| centers, scales = self._get_centers_and_scales_xyxy([m["bbox"] for m in meta]) |
| output = self.model(tensor).detach().cpu().numpy() |
| kps, scores = self.keypoints_from_heatmaps(output, centers, scales, unbiased=True, post_process="default", target_type="GaussianHeatmap", kernel=11) |
| return kps, scores |
|
|
| @staticmethod |
| def visualize(image_bgr, keypoints, scores, score_threshold=0.3, kp_radius=8, line_thickness=5): |
| canvas = image_bgr.copy() |
| if keypoints.ndim == 2: |
| keypoints = np.expand_dims(keypoints, axis=0) |
| scores = np.expand_dims(scores, axis=0) |
| edges = [(0,1),(0,2),(1,3),(2,4),(5,6),(5,11),(6,12),(11,12),(5,7),(7,9),(6,8),(8,10),(11,13),(13,15),(12,14),(14,16)] |
| colors = [(255,0,0),(255,85,0),(255,170,0),(255,255,0),(170,255,0),(85,255,0),(0,255,0),(0,255,85),(0,255,170),(0,255,255),(0,170,255),(0,85,255),(0,0,255),(85,0,255),(170,0,255),(255,0,255)] |
| for n in range(len(keypoints)): |
| kpts, scs = keypoints[n], scores[n].squeeze() |
| for i,(a,b) in enumerate(edges): |
| if scs[a]>=score_threshold and scs[b]>=score_threshold: |
| cv2.line(canvas,(int(kpts[a][0]),int(kpts[a][1])),(int(kpts[b][0]),int(kpts[b][1])),colors[i],thickness=line_thickness) |
| for k in range(len(kpts)): |
| if scs[k]>=score_threshold: |
| cv2.circle(canvas,(int(kpts[k,0]),int(kpts[k,1])),kp_radius,color=(255,255,255),thickness=-1) |
| return canvas |
|
|
