Added image output to OpenPoseToPointList

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)