mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-02-02 11:55:17 -05:00
113 lines
3.3 KiB
Python
113 lines
3.3 KiB
Python
from pathlib import Path
|
|
from typing import Dict
|
|
|
|
import numpy as np
|
|
import torch
|
|
from controlnet_aux.util import resize_image
|
|
from PIL import Image
|
|
|
|
from invokeai.backend.image_util.dw_openpose.utils import NDArrayInt, draw_bodypose, draw_facepose, draw_handpose
|
|
from invokeai.backend.image_util.dw_openpose.wholebody import Wholebody
|
|
|
|
DWPOSE_MODELS = {
|
|
"yolox_l.onnx": "https://huggingface.co/yzd-v/DWPose/resolve/main/yolox_l.onnx?download=true",
|
|
"dw-ll_ucoco_384.onnx": "https://huggingface.co/yzd-v/DWPose/resolve/main/dw-ll_ucoco_384.onnx?download=true",
|
|
}
|
|
|
|
|
|
def draw_pose(
|
|
pose: Dict[str, NDArrayInt | Dict[str, NDArrayInt]],
|
|
H: int,
|
|
W: int,
|
|
draw_face: bool = True,
|
|
draw_body: bool = True,
|
|
draw_hands: bool = True,
|
|
resolution: int = 512,
|
|
) -> Image.Image:
|
|
bodies = pose["bodies"]
|
|
faces = pose["faces"]
|
|
hands = pose["hands"]
|
|
|
|
assert isinstance(bodies, dict)
|
|
candidate = bodies["candidate"]
|
|
|
|
assert isinstance(bodies, dict)
|
|
subset = bodies["subset"]
|
|
|
|
canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
|
|
|
|
if draw_body:
|
|
canvas = draw_bodypose(canvas, candidate, subset)
|
|
|
|
if draw_hands:
|
|
assert isinstance(hands, np.ndarray)
|
|
canvas = draw_handpose(canvas, hands)
|
|
|
|
if draw_face:
|
|
assert isinstance(hands, np.ndarray)
|
|
canvas = draw_facepose(canvas, faces) # type: ignore
|
|
|
|
dwpose_image: Image.Image = resize_image(
|
|
canvas,
|
|
resolution,
|
|
)
|
|
dwpose_image = Image.fromarray(dwpose_image)
|
|
|
|
return dwpose_image
|
|
|
|
|
|
class DWOpenposeDetector:
|
|
"""
|
|
Code from the original implementation of the DW Openpose Detector.
|
|
Credits: https://github.com/IDEA-Research/DWPose
|
|
"""
|
|
|
|
def __init__(self, onnx_det: Path, onnx_pose: Path) -> None:
|
|
self.pose_estimation = Wholebody(onnx_det=onnx_det, onnx_pose=onnx_pose)
|
|
|
|
def __call__(
|
|
self,
|
|
image: Image.Image,
|
|
draw_face: bool = False,
|
|
draw_body: bool = True,
|
|
draw_hands: bool = False,
|
|
resolution: int = 512,
|
|
) -> Image.Image:
|
|
np_image = np.array(image)
|
|
H, W, C = np_image.shape
|
|
|
|
with torch.no_grad():
|
|
candidate, subset = self.pose_estimation(np_image)
|
|
nums, keys, locs = candidate.shape
|
|
candidate[..., 0] /= float(W)
|
|
candidate[..., 1] /= float(H)
|
|
body = candidate[:, :18].copy()
|
|
body = body.reshape(nums * 18, locs)
|
|
score = subset[:, :18]
|
|
for i in range(len(score)):
|
|
for j in range(len(score[i])):
|
|
if score[i][j] > 0.3:
|
|
score[i][j] = int(18 * i + j)
|
|
else:
|
|
score[i][j] = -1
|
|
|
|
un_visible = subset < 0.3
|
|
candidate[un_visible] = -1
|
|
|
|
# foot = candidate[:, 18:24]
|
|
|
|
faces = candidate[:, 24:92]
|
|
|
|
hands = candidate[:, 92:113]
|
|
hands = np.vstack([hands, candidate[:, 113:]])
|
|
|
|
bodies = {"candidate": body, "subset": score}
|
|
pose = {"bodies": bodies, "hands": hands, "faces": faces}
|
|
|
|
return draw_pose(
|
|
pose, H, W, draw_face=draw_face, draw_hands=draw_hands, draw_body=draw_body, resolution=resolution
|
|
)
|
|
|
|
|
|
__all__ = ["DWPOSE_MODELS", "DWOpenposeDetector"]
|