# Openpose
# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
# 3rd Edited by ControlNet
# 4th Edited by ControlNet (added face and correct hands)
# 5th Edited by ControlNet (Improved JSON serialization/deserialization, and lots of bug fixs)
# This preprocessor is licensed by CMU for non-commercial use only.
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
import json
import torch
import numpy as np
from . import util
from .body import Body, BodyResult, Keypoint
from .hand import Hand
from .face import Face
from .types import PoseResult, HandResult, FaceResult
from annotator.annotator_path import models_path
from typing import Tuple, List, Callable, Union, Optional
body_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/body_pose_model.pth"
hand_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/hand_pose_model.pth"
face_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/facenet.pth"
remote_onnx_det = "https://huggingface.co/yzd-v/DWPose/resolve/main/yolox_l.onnx"
remote_onnx_pose = "https://huggingface.co/yzd-v/DWPose/resolve/main/dw-ll_ucoco_384.onnx"
def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, draw_face=True):
"""
Draw the detected poses on an empty canvas.
Args:
poses (List[PoseResult]): A list of PoseResult objects containing the detected poses.
H (int): The height of the canvas.
W (int): The width of the canvas.
draw_body (bool, optional): Whether to draw body keypoints. Defaults to True.
draw_hand (bool, optional): Whether to draw hand keypoints. Defaults to True.
draw_face (bool, optional): Whether to draw face keypoints. Defaults to True.
Returns:
numpy.ndarray: A 3D numpy array representing the canvas with the drawn poses.
"""
canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
for pose in poses:
if draw_body:
canvas = util.draw_bodypose(canvas, pose.body.keypoints)
if draw_hand:
canvas = util.draw_handpose(canvas, pose.left_hand)
canvas = util.draw_handpose(canvas, pose.right_hand)
if draw_face:
canvas = util.draw_facepose(canvas, pose.face)
return canvas
def decode_json_as_poses(json_string: str, normalize_coords: bool = False) -> Tuple[List[PoseResult], int, int]:
""" Decode the json_string complying with the openpose JSON output format
to poses that controlnet recognizes.
https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/02_output.md
Args:
json_string: The json string to decode.
normalize_coords: Whether to normalize coordinates of each keypoint by canvas height/width.
`draw_pose` only accepts normalized keypoints. Set this param to True if
the input coords are not normalized.
Returns:
poses
canvas_height
canvas_width
"""
pose_json = json.loads(json_string)
height = pose_json['canvas_height']
width = pose_json['canvas_width']
def chunks(lst, n):
"""Yield successive n-sized chunks from lst."""
for i in range(0, len(lst), n):
yield lst[i:i + n]
def decompress_keypoints(numbers: Optional[List[float]]) -> Optional[List[Optional[Keypoint]]]:
if not numbers:
return None
assert len(numbers) % 3 == 0
def create_keypoint(x, y, c):
if c < 1.0:
return None
keypoint = Keypoint(x, y)
return keypoint
return [
create_keypoint(x, y, c)
for x, y, c in chunks(numbers, n=3)
]
return (
[
PoseResult(
body=BodyResult(keypoints=decompress_keypoints(pose.get('pose_keypoints_2d'))),
left_hand=decompress_keypoints(pose.get('hand_left_keypoints_2d')),
right_hand=decompress_keypoints(pose.get('hand_right_keypoints_2d')),
face=decompress_keypoints(pose.get('face_keypoints_2d'))
)
for pose in pose_json['people']
],
height,
width,
)
def encode_poses_as_json(poses: List[PoseResult], canvas_height: int, canvas_width: int) -> dict:
""" Encode the pose as a JSON compatible dict following openpose JSON output format:
https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/02_output.md
"""
def compress_keypoints(keypoints: Union[List[Keypoint], None]) -> Union[List[float], None]:
if not keypoints:
return None
return [
value
for keypoint in keypoints
for value in (
[float(keypoint.x), float(keypoint.y), 1.0]
if keypoint is not None
else [0.0, 0.0, 0.0]
)
]
return {
'people': [
{
'pose_keypoints_2d': compress_keypoints(pose.body.keypoints),
"face_keypoints_2d": compress_keypoints(pose.face),
"hand_left_keypoints_2d": compress_keypoints(pose.left_hand),
"hand_right_keypoints_2d":compress_keypoints(pose.right_hand),
}
for pose in poses
],
'canvas_height': canvas_height,
'canvas_width': canvas_width,
}
class OpenposeDetector:
"""
A class for detecting human poses in images using the Openpose model.
Attributes:
model_dir (str): Path to the directory where the pose models are stored.
"""
model_dir = os.path.join(models_path, "openpose")
def __init__(self):
self.device = 'cuda'
self.body_estimation = None
self.hand_estimation = None
self.face_estimation = None
self.dw_pose_estimation = None
def load_model(self):
"""
Load the Openpose body, hand, and face models.
"""
body_modelpath = os.path.join(self.model_dir, "body_pose_model.pth")
hand_modelpath = os.path.join(self.model_dir, "hand_pose_model.pth")
face_modelpath = os.path.join(self.model_dir, "facenet.pth")
if not os.path.exists(body_modelpath):
from basicsr.utils.download_util import load_file_from_url
load_file_from_url(body_model_path, model_dir=self.model_dir)
if not os.path.exists(hand_modelpath):
from basicsr.utils.download_util import load_file_from_url
load_file_from_url(hand_model_path, model_dir=self.model_dir)
if not os.path.exists(face_modelpath):
from basicsr.utils.download_util import load_file_from_url
load_file_from_url(face_model_path, model_dir=self.model_dir)
self.body_estimation = Body(body_modelpath)
self.hand_estimation = Hand(hand_modelpath)
self.face_estimation = Face(face_modelpath)
def load_dw_model(self):
from .wholebody import Wholebody # DW Pose
def load_model(filename: str, remote_url: str):
local_path = os.path.join(self.model_dir, filename)
if not os.path.exists(local_path):
from basicsr.utils.download_util import load_file_from_url
load_file_from_url(remote_url, model_dir=self.model_dir)
return local_path
onnx_det = load_model("yolox_l.onnx", remote_onnx_det)
onnx_pose = load_model("dw-ll_ucoco_384.onnx", remote_onnx_pose)
self.dw_pose_estimation = Wholebody(onnx_det, onnx_pose)
def unload_model(self):
"""
Unload the Openpose models by moving them to the CPU.
Note: DW Pose models always run on CPU, so no need to `unload` them.
"""
if self.body_estimation is not None:
self.body_estimation.model.to("cpu")
self.hand_estimation.model.to("cpu")
self.face_estimation.model.to("cpu")
def detect_hands(self, body: BodyResult, oriImg) -> Tuple[Union[HandResult, None], Union[HandResult, None]]:
left_hand = None
right_hand = None
H, W, _ = oriImg.shape
for x, y, w, is_left in util.handDetect(body, oriImg):
peaks = self.hand_estimation(oriImg[y:y+w, x:x+w, :]).astype(np.float32)
if peaks.ndim == 2 and peaks.shape[1] == 2:
peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)
peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)
hand_result = [
Keypoint(x=peak[0], y=peak[1])
for peak in peaks
]
if is_left:
left_hand = hand_result
else:
right_hand = hand_result
return left_hand, right_hand
def detect_face(self, body: BodyResult, oriImg) -> Union[FaceResult, None]:
face = util.faceDetect(body, oriImg)
if face is None:
return None
x, y, w = face
H, W, _ = oriImg.shape
heatmaps = self.face_estimation(oriImg[y:y+w, x:x+w, :])
peaks = self.face_estimation.compute_peaks_from_heatmaps(heatmaps).astype(np.float32)
if peaks.ndim == 2 and peaks.shape[1] == 2:
peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)
peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)
return [
Keypoint(x=peak[0], y=peak[1])
for peak in peaks
]
return None
def detect_poses(self, oriImg, include_hand=False, include_face=False) -> List[PoseResult]:
"""
Detect poses in the given image.
Args:
oriImg (numpy.ndarray): The input image for pose detection.
include_hand (bool, optional): Whether to include hand detection. Defaults to False.
include_face (bool, optional): Whether to include face detection. Defaults to False.
Returns:
List[PoseResult]: A list of PoseResult objects containing the detected poses.
"""
if self.body_estimation is None:
self.load_model()
self.body_estimation.model.to(self.device)
self.hand_estimation.model.to(self.device)
self.face_estimation.model.to(self.device)
self.body_estimation.cn_device = self.device
self.hand_estimation.cn_device = self.device
self.face_estimation.cn_device = self.device
oriImg = oriImg[:, :, ::-1].copy()
H, W, C = oriImg.shape
with torch.no_grad():
candidate, subset = self.body_estimation(oriImg)
bodies = self.body_estimation.format_body_result(candidate, subset)
results = []
for body in bodies:
left_hand, right_hand, face = (None,) * 3
if include_hand:
left_hand, right_hand = self.detect_hands(body, oriImg)
if include_face:
face = self.detect_face(body, oriImg)
results.append(PoseResult(BodyResult(
keypoints=[
Keypoint(
x=keypoint.x / float(W),
y=keypoint.y / float(H)
) if keypoint is not None else None
for keypoint in body.keypoints
],
total_score=body.total_score,
total_parts=body.total_parts
), left_hand, right_hand, face))
return results
def detect_poses_dw(self, oriImg) -> List[PoseResult]:
"""
Detect poses in the given image using DW Pose:
https://github.com/IDEA-Research/DWPose
Args:
oriImg (numpy.ndarray): The input image for pose detection.
Returns:
List[PoseResult]: A list of PoseResult objects containing the detected poses.
"""
from .wholebody import Wholebody # DW Pose
self.load_dw_model()
with torch.no_grad():
keypoints_info = self.dw_pose_estimation(oriImg.copy())
return Wholebody.format_result(keypoints_info)
def __call__(
self, oriImg, include_body=True, include_hand=False, include_face=False,
use_dw_pose=False, json_pose_callback: Callable[[str], None] = None,
):
"""
Detect and draw poses in the given image.
Args:
oriImg (numpy.ndarray): The input image for pose detection and drawing.
include_body (bool, optional): Whether to include body keypoints. Defaults to True.
include_hand (bool, optional): Whether to include hand keypoints. Defaults to False.
include_face (bool, optional): Whether to include face keypoints. Defaults to False.
use_dw_pose (bool, optional): Whether to use DW pose detection algorithm. Defaults to False.
json_pose_callback (Callable, optional): A callback that accepts the pose JSON string.
Returns:
numpy.ndarray: The image with detected and drawn poses.
"""
H, W, _ = oriImg.shape
if use_dw_pose:
poses = self.detect_poses_dw(oriImg)
else:
poses = self.detect_poses(oriImg, include_hand, include_face)
if json_pose_callback:
json_pose_callback(encode_poses_as_json(poses, H, W))
return draw_poses(poses, H, W, draw_body=include_body, draw_hand=include_hand, draw_face=include_face)