acmyu/KeyframesAI

1from typing import Callable, Dict, Optional, Union2 3import cv24import numpy as np5import PIL6import PIL.Image7import torch8from huggingface_hub import hf_hub_download9 10from .body_estimation import Wholebody, resize_image11from .draw import draw_openpose12 13 14class DWposeDetector:15    def __init__(self, device: str = "сpu"):16        hf_hub_download("RedHash/DWPose", "yolox_l.onnx", local_dir="./checkpoints")17        hf_hub_download("RedHash/DWPose", "dw-ll_ucoco_384.onnx", local_dir="./checkpoints")18        self.pose_estimation = Wholebody(19            device=device, model_det="checkpoints/yolox_l.onnx", model_pose="checkpoints/dw-ll_ucoco_384.onnx"20        )21 22    def _format_pose(self, candidates, scores, width, height):23        num_candidates, _, locs = candidates.shape24 25        candidates[..., 0] /= float(width)26        candidates[..., 1] /= float(height)27 28        bodies = candidates[:, :18].copy()29        bodies_flat = bodies.reshape(num_candidates * 18, locs)30 31        body_scores = scores[:, :18]32        for i in range(len(body_scores)):33            for j in range(len(body_scores[i])):34                if body_scores[i][j] > 0.3:35                    body_scores[i][j] = int(18 * i + j)36                else:37                    body_scores[i][j] = -138 39        faces = candidates[:, 24:92]40        faces_scores = scores[:, 24:92]41 42        hands = np.vstack([candidates[:, 92:113], candidates[:, 113:]])43        hands_scores = np.vstack([scores[:, 92:113], scores[:, 113:]])44 45        pose = dict(46            bodies=bodies_flat,47            bodies_multi=bodies,48            body_scores=body_scores,49            hands=hands,50            hands_scores=hands_scores,51            faces=faces,52            faces_scores=faces_scores,53            num_candidates=num_candidates,54        )55 56        return pose57 58    @torch.inference_mode()59    def __call__(60        self,61        image: Union[PIL.Image.Image, np.ndarray],62        detect_resolution: int = 512,63        draw_pose: Optional[Callable] = draw_openpose,64        output_type: str = "pil",65        **kwargs,66    ) -> Union[PIL.Image.Image, np.ndarray, Dict]:67        if type(image) != np.ndarray:68            image = np.array(image.convert("RGB"))69 70        image = image.copy()71        original_height, original_width, _ = image.shape72 73        image = resize_image(image, target_resolution=detect_resolution)74        height, width, _ = image.shape75 76        candidates, scores = self.pose_estimation(image)77 78        pose = self._format_pose(candidates, scores, width, height)79 80        if not draw_pose:81            return pose82 83        pose_image = draw_pose(pose, height=height, width=width, **kwargs)84        pose_image = cv2.resize(pose_image, (original_width, original_height), cv2.INTER_LANCZOS4)85 86        if output_type == "pil":87            pose_image = PIL.Image.fromarray(pose_image)88        elif output_type == "np":89            pass90        else:91            raise ValueError("output_type should be 'pil' or 'np'")92 93        return pose_image, pose94