| from __future__ import annotations |
|
|
| import os |
| from typing import Optional, Sequence, Tuple, List |
|
|
| import torch |
| import torch.distributed as dist |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch import Tensor |
| from collections import OrderedDict |
| import warnings |
| import logging |
| import numpy as np |
| import cv2 |
| from safetensors.torch import load_file |
| from huggingface_hub import hf_hub_download |
|
|
| def _drop_path(x: Tensor, drop_prob: float = 0., training: bool = False) -> Tensor: |
| if drop_prob == 0. or not training: |
| return x |
| keep_prob = 1 - drop_prob |
| shape = (x.shape[0],) + (1,) * (x.ndim - 1) |
| random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) |
| random_tensor.floor_() |
| output = x.div(keep_prob) * random_tensor |
| return output |
|
|
| class _DropPath(nn.Module): |
| def __init__(self, drop_prob: float = 0.) -> None: |
| super().__init__() |
| self.drop_prob = drop_prob |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| return _drop_path(x, self.drop_prob, self.training) |
|
|
| class _PConv2d(nn.Module): |
| def __init__(self, in_channels: int, kernel_size: int = 3, n_div: int = 4) -> None: |
| super().__init__() |
| self.dim_conv = in_channels // n_div |
| self.dim_untouched = in_channels - self.dim_conv |
| self.conv = nn.Conv2d( |
| self.dim_conv, self.dim_conv, kernel_size, stride=1, |
| padding=(kernel_size - 1) // 2, bias=False |
| ) |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| x1, x2 = torch.split(x, [self.dim_conv, self.dim_untouched], dim=1) |
| x1 = self.conv(x1) |
| return torch.cat((x1, x2), dim=1) |
|
|
| class _ConvBNLayer(nn.Module): |
| def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 1, |
| stride: int = 1, padding: int = 0, act: str = 'ReLU') -> None: |
| super().__init__() |
| self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False) |
| self.bn = nn.BatchNorm2d(out_channels) |
| self.act = getattr(nn, act)() |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| return self.act(self.bn(self.conv(x))) |
|
|
| class _FasterNetBlock(nn.Module): |
| def __init__(self, in_channels: int, inner_channels: int = None, kernel_size: int = 3, |
| act: str = 'ReLU', n_div: int = 4, drop_path: float = 0.) -> None: |
| super().__init__() |
| inner_channels = inner_channels or in_channels * 2 |
| self.conv1 = _PConv2d(in_channels, kernel_size, n_div) |
| self.conv2 = _ConvBNLayer(in_channels, inner_channels, act=act) |
| self.conv3 = nn.Conv2d(inner_channels, in_channels, kernel_size=1, stride=1, bias=False) |
| self.drop_path = _DropPath(drop_path) if drop_path > 0. else nn.Identity() |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| y = self.conv1(x) |
| y = self.conv2(y) |
| y = self.conv3(y) |
| return x + self.drop_path(y) |
|
|
| class _FasterNet(nn.Module): |
| def __init__(self, in_channels: int = 3, inner_channels: List[int] = [192, 384, 768, 1536], |
| blocks: List[int] = [3, 4, 18, 3], act: str = 'ReLU', drop_path: float = 0.3) -> None: |
| super().__init__() |
| self.embedding = _ConvBNLayer(in_channels, inner_channels[0], kernel_size=4, stride=4) |
|
|
| self.stage1 = nn.Sequential(OrderedDict([ |
| (f'block{i}', _FasterNetBlock(inner_channels[0], act=act, drop_path=drop_path)) |
| for i in range(blocks[0]) |
| ])) |
|
|
| self.merging1 = _ConvBNLayer(inner_channels[0], inner_channels[1], kernel_size=2, stride=2) |
| self.stage2 = nn.Sequential(OrderedDict([ |
| (f'block{i}', _FasterNetBlock(inner_channels[1], act=act, drop_path=drop_path)) |
| for i in range(blocks[1]) |
| ])) |
|
|
| self.merging2 = _ConvBNLayer(inner_channels[1], inner_channels[2], kernel_size=2, stride=2) |
| self.stage3 = nn.Sequential(OrderedDict([ |
| (f'block{i}', _FasterNetBlock(inner_channels[2], act=act, drop_path=drop_path)) |
| for i in range(blocks[2]) |
| ])) |
|
|
| self.merging3 = _ConvBNLayer(inner_channels[2], inner_channels[3], kernel_size=2, stride=2) |
| self.stage4 = nn.Sequential(OrderedDict([ |
| (f'block{i}', _FasterNetBlock(inner_channels[3], act=act, drop_path=drop_path)) |
| for i in range(blocks[3]) |
| ])) |
|
|
| def forward(self, x: Tensor) -> List[Tensor]: |
| x1 = self.stage1(self.embedding(x)) |
| x2 = self.stage2(self.merging1(x1)) |
| x3 = self.stage3(self.merging2(x2)) |
| x4 = self.stage4(self.merging3(x3)) |
| return [x1, x2, x3, x4] |
|
|
| class _FasterNetBackbone(nn.Module): |
| def __init__( |
| self, |
| variant: str = 's', |
| frozen_stages: int = -1, |
| norm_eval: bool = False, |
| ) -> None: |
| super().__init__() |
|
|
| self.variant = variant |
| self.frozen_stages = frozen_stages |
| self.norm_eval = norm_eval |
|
|
| if variant == 's': |
| _net = _FasterNet(inner_channels=[128, 256, 512, 1024], blocks=[1, 2, 13, 2], drop_path=0.1) |
| self.out_channels = 1024 |
| elif variant == 'm': |
| _net = _FasterNet(inner_channels=[144, 288, 576, 1152], blocks=[3, 4, 18, 3], drop_path=0.2) |
| self.out_channels = 1152 |
| elif variant == 'l': |
| _net = _FasterNet(inner_channels=[192, 384, 768, 1536], blocks=[3, 4, 18, 3], drop_path=0.3) |
| self.out_channels = 1536 |
| else: |
| raise ValueError(f"Expected variant 's', 'm', or 'l', got {variant}") |
| self.embedding = _net.embedding |
| self.stage1 = _net.stage1 |
| self.merging1 = _net.merging1 |
| self.stage2 = _net.stage2 |
| self.merging2 = _net.merging2 |
| self.stage3 = _net.stage3 |
| self.merging3 = _net.merging3 |
| self.stage4 = _net.stage4 |
|
|
|
|
| self.classifier = nn.Sequential( |
| OrderedDict([ |
| ('conv', nn.Conv2d(self.out_channels, 1280, kernel_size=1, bias=False)), |
| ('fc', nn.Linear(1280, 1000)) |
| ]) |
| ) |
| self._freeze_stages() |
|
|
| def _freeze_stages(self) -> None: |
| stages = [self.embedding, self.stage1, self.merging1, |
| self.stage2, self.merging2, self.stage3, |
| self.merging3, self.stage4] |
|
|
| for i, stage in enumerate(stages): |
| if i <= self.frozen_stages: |
| stage.eval() |
| for param in stage.parameters(): |
| param.requires_grad = False |
|
|
| def train(self, mode: bool = True) -> 'FasterNetBackbone': |
| 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, ...]: |
| x1 = self.stage1(self.embedding(x)) |
| x2 = self.stage2(self.merging1(x1)) |
| x3 = self.stage3(self.merging2(x2)) |
| x4 = self.stage4(self.merging3(x3)) |
| return (x1, x2, x3, x4) |
|
|
| 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 FasterNetSPoseNet(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 = _FasterNetBackbone( |
| variant='s', |
| 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 = FasterNetSPoseNet(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 |
|
|
