"""Singleton predictor for ViTPose human pose estimation.""" import logging import math import cv2 import numpy as np import torch from numpy.typing import NDArray from PIL import Image from transformers import VitPoseForPoseEstimation from transformers.models.vitpose.image_processing_vitpose import VitPoseImageProcessor from .model import ( KEYPOINT_COLOR_INDICES, KEYPOINT_RADIUS, KEYPOINT_SCORE_THRESHOLD, LIMB_THICKNESS, LINK_COLOR_INDICES, PALETTE, load_model, ) logger = logging.getLogger(__name__) def _draw_keypoints( image: NDArray[np.uint8], keypoints: NDArray[np.float64], scores: NDArray[np.float64], keypoint_colors: NDArray[np.int64], radius: int, threshold: float, ) -> None: """Draw keypoint circles on the image. Args: image: BGR image array (H, W, 3), modified in-place. keypoints: Keypoint coordinates of shape (num_keypoints, 2). scores: Confidence scores of shape (num_keypoints,). keypoint_colors: RGB color array of shape (num_keypoints, 3). radius: Circle radius in pixels. threshold: Minimum score to draw a keypoint. """ for kid, (kpt, kpt_score) in enumerate(zip(keypoints, scores, strict=True)): if kpt_score <= threshold: continue x_coord, y_coord = int(kpt[0]), int(kpt[1]) color = tuple(int(c) for c in keypoint_colors[kid]) cv2.circle(image, (x_coord, y_coord), radius, color, -1) def _draw_limbs( image: NDArray[np.uint8], keypoints: NDArray[np.float64], scores: NDArray[np.float64], edges: list[tuple[int, int]], link_colors: NDArray[np.int64], thickness: int, threshold: float, ) -> None: """Draw skeleton limb connections on the image. Args: image: BGR image array (H, W, 3), modified in-place. keypoints: Keypoint coordinates of shape (num_keypoints, 2). scores: Confidence scores of shape (num_keypoints,). edges: List of (start_idx, end_idx) pairs defining limb connections. link_colors: RGB color array of shape (num_edges, 3). thickness: Line thickness in pixels. threshold: Minimum score for both endpoints to draw a limb. """ height, width, _ = image.shape for sk_id, (start_idx, end_idx) in enumerate(edges): x1, y1 = int(keypoints[start_idx, 0]), int(keypoints[start_idx, 1]) x2, y2 = int(keypoints[end_idx, 0]), int(keypoints[end_idx, 1]) score1 = scores[start_idx] score2 = scores[end_idx] in_bounds = ( 0 < x1 < width and 0 < y1 < height and 0 < x2 < width and 0 < y2 < height ) above_threshold = score1 > threshold and score2 > threshold if not (in_bounds and above_threshold): continue color = tuple(int(c) for c in link_colors[sk_id]) mean_x = int(np.mean([x1, x2])) mean_y = int(np.mean([y1, y2])) length = int(((y1 - y2) ** 2 + (x1 - x2) ** 2) ** 0.5) angle = int(math.degrees(math.atan2(y1 - y2, x1 - x2))) stick_width = 2 polygon = cv2.ellipse2Poly( (mean_x, mean_y), (length // 2, stick_width), angle, 0, 360, 1, ) cv2.fillConvexPoly(image, polygon, color) class Predictor: """Singleton predictor for ViTPose pose estimation. Usage: predictor = Predictor() predictor.initialize(device="cuda") result_image = predictor.predict(input_image) """ _instance: "Predictor | None" = None _initialized: bool = False def __new__(cls) -> "Predictor": if cls._instance is None: cls._instance = super().__new__(cls) return cls._instance def __init__(self) -> None: if not hasattr(self, "_model"): self._model: VitPoseForPoseEstimation | None = None self._processor: VitPoseImageProcessor | None = None self._device: str = "cpu" def initialize(self, device: str = "cuda") -> None: """Load the model and processor. Args: device: Device string ("cuda" or "cpu"). """ if self._initialized: logger.info("Predictor already initialized, skipping") return self._device = device self._model, self._processor = load_model(device) self._initialized = True logger.info("Predictor initialized on %s", self._device) @property def model(self) -> VitPoseForPoseEstimation: """Return the loaded model, raising if not initialized.""" if self._model is None: raise RuntimeError("Predictor not initialized. Call initialize() first.") return self._model @property def processor(self) -> VitPoseImageProcessor: """Return the loaded processor, raising if not initialized.""" if self._processor is None: raise RuntimeError("Predictor not initialized. Call initialize() first.") return self._processor def predict(self, image: Image.Image) -> Image.Image: """Detect keypoints and draw skeleton overlay on the image. Uses the full image as a single bounding box (single-person mode). For multi-person detection, an object detector would be needed upstream. Args: image: Input PIL image. Returns: PIL image with skeleton overlay drawn. """ if not self._initialized: raise RuntimeError("Predictor not initialized. Call initialize() first.") image_rgb = image.convert("RGB") width, height = image_rgb.size # Use the full image as a single bounding box (COCO format: x, y, w, h) boxes = [[[0.0, 0.0, float(width), float(height)]]] inputs = self.processor(image_rgb, boxes=boxes, return_tensors="pt").to( self.model.device ) with torch.no_grad(): outputs = self.model(**inputs) pose_results = self.processor.post_process_pose_estimation(outputs, boxes=boxes) image_pose_result = pose_results[0] # Get skeleton edges from model config edges: list[tuple[int, int]] = self.model.config.edges # Prepare color arrays palette = np.array(PALETTE) link_colors = palette[LINK_COLOR_INDICES] keypoint_colors = palette[KEYPOINT_COLOR_INDICES] # Draw on a copy of the image (OpenCV uses BGR) numpy_image = np.array(image_rgb) canvas = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR) for pose_result in image_pose_result: keypoints = np.array(pose_result["keypoints"]) scores = np.array(pose_result["scores"]) _draw_limbs( canvas, keypoints, scores, edges, link_colors, thickness=LIMB_THICKNESS, threshold=KEYPOINT_SCORE_THRESHOLD, ) _draw_keypoints( canvas, keypoints, scores, keypoint_colors, radius=KEYPOINT_RADIUS, threshold=KEYPOINT_SCORE_THRESHOLD, ) # Convert back to RGB for PIL result_rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB) return Image.fromarray(result_rgb)