diff --git a/negi/open_pose_to_point_list.py b/negi/open_pose_to_point_list.py index 75a0293..3b0f58a 100644 --- a/negi/open_pose_to_point_list.py +++ b/negi/open_pose_to_point_list.py @@ -2,8 +2,10 @@ import json import numpy as np import cv2 +import torch from controlnet_aux import OpenposeDetector from controlnet_aux.util import HWC3 +from controlnet_aux.open_pose import draw_poses _names = [ @@ -28,7 +30,7 @@ def _resize_image(input_image, resolution): 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 + return img, H, W class OpenPoseToPointList: @@ -49,8 +51,8 @@ class OpenPoseToPointList: }, } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("POINT_LIST",) + RETURN_TYPES = ("STRING", "IMAGE") + RETURN_NAMES = ("POINT_LIST", "IMAGE") FUNCTION = "doit" OUTPUT_NODE = False CATEGORY = "utils" @@ -58,9 +60,11 @@ class OpenPoseToPointList: 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 = _resize_image(input_image, detect_resolution) + 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 = [] @@ -97,4 +101,4 @@ class OpenPoseToPointList: else: raise ValueError() - return (json.dumps(ret, indent=2),) + return (json.dumps(ret, indent=2), img)