105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
import json
|
|
|
|
import numpy as np
|
|
import cv2
|
|
import torch
|
|
from .repos.controlnet_aux.src.controlnet_aux import OpenposeDetector
|
|
from .repos.controlnet_aux.src.controlnet_aux.util import HWC3
|
|
from .repos.controlnet_aux.src.controlnet_aux.open_pose import draw_poses
|
|
|
|
|
|
_names = [
|
|
"Nose", "Neck",
|
|
"RShoulder", "RElbow", "RWrist",
|
|
"LShoulder", "LElbow", "LWrist",
|
|
"RHip", "RKnee", "RAnkle",
|
|
"LHip", "LKnee", "LAnkle",
|
|
"REye", "LEye", "REar", "LEar"
|
|
]
|
|
|
|
_name_to_index = {name: i for i, name in enumerate(_names)}
|
|
|
|
|
|
def _resize_image(input_image, resolution):
|
|
H, W, C = input_image.shape
|
|
H = float(H)
|
|
W = float(W)
|
|
k = float(resolution) / min(H, W)
|
|
H *= k
|
|
W *= k
|
|
H = int(np.round(H / 64.0)) * 64
|
|
W = int(np.round(W / 64.0)) * 64
|
|
img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
|
|
return img, H, W
|
|
|
|
|
|
class OpenPoseToPointList:
|
|
def __init__(self):
|
|
self.open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"detect_resolution": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64, "display": "slider"}),
|
|
"method": ([
|
|
"face",
|
|
"hand",
|
|
"all",
|
|
],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "IMAGE")
|
|
RETURN_NAMES = ("POINT_LIST", "IMAGE")
|
|
FUNCTION = "doit"
|
|
OUTPUT_NODE = False
|
|
CATEGORY = "utils"
|
|
|
|
def doit(self, image, detect_resolution, method):
|
|
input_image = (np.fmax(0.0, np.fmin(1.0, image.to('cpu').detach().numpy()[0])) * 255.0).astype(np.uint8)
|
|
input_image = HWC3(input_image)
|
|
input_image, H, W = _resize_image(input_image, detect_resolution)
|
|
|
|
poses = self.open_pose.detect_poses(input_image, include_hand=False, include_face=False)
|
|
img = draw_poses(poses, H, W, draw_hand=False, draw_face=False)
|
|
img = torch.from_numpy(np.expand_dims(HWC3(img) * (1.0 / 255), axis=0))
|
|
|
|
if method == "face":
|
|
ret = []
|
|
for pose in poses:
|
|
x = 0.0
|
|
y = 0.0
|
|
n = 0
|
|
for name in ["Nose", "REye", "LEye", "REar", "LEar"]:
|
|
key_point = pose.body.keypoints[_name_to_index[name]]
|
|
if key_point is not None:
|
|
x += key_point.x
|
|
y += key_point.y
|
|
n += 1
|
|
if n != 0:
|
|
ret.append({"x": x / n, "y": y / n})
|
|
|
|
elif method == "hand":
|
|
ret = []
|
|
for pose in poses:
|
|
for name in ["RWrist", "LWrist"]:
|
|
key_point = pose.body.keypoints[_name_to_index[name]]
|
|
if key_point is not None:
|
|
ret.append({"x": key_point.x, "y": key_point.y})
|
|
|
|
elif method == "all":
|
|
ret = []
|
|
for pose in poses:
|
|
points = {}
|
|
for i, key_point in enumerate(pose.body.keypoints):
|
|
if key_point is not None:
|
|
points[_names[i]] = {"x": key_point.x, "y": key_point.y, "score": key_point.score}
|
|
ret.append(points)
|
|
|
|
else:
|
|
raise ValueError()
|
|
|
|
return (json.dumps(ret, indent=2), img)
|