From 9075ec965b0d9a0d77b3ce2d127625bdfd3c2121 Mon Sep 17 00:00:00 2001
From: NiggoJaecha <andre.grabowich@gmail.com>
Date: Wed, 22 Feb 2023 23:04:19 +0100
Subject: [PATCH] Fix Openpose and Segmentation

Fix Openpose: Add code explicitly for openpose
Segmentations (Uniformer): Fix path mismatch (config is not in models folder)
---
 custom_nodes/annotator/__init__.py           | 15 +++++++++------
 custom_nodes/annotator/uniformer/__init__.py |  2 +-
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/custom_nodes/annotator/__init__.py b/custom_nodes/annotator/__init__.py
index 9edda1c49..ecb01650b 100644
--- a/custom_nodes/annotator/__init__.py
+++ b/custom_nodes/annotator/__init__.py
@@ -14,12 +14,15 @@ def img_tensor_to_np(img_tensor):
 
 def common_annotator_call(annotator_callback, tensor_image, *args):
     call_result = annotator_callback(img_tensor_to_np(tensor_image), *args)
-    if type(call_result) is tuple:
-        for i in range(len(call_result)):
-            call_result[i] = HWC3(call_result[i])
-    else:
-        call_result = HWC3(call_result)
-    return call_result
+    if type(annotator_callback) is openpose.OpenposeDetector:
+        return (HWC3(call_result[0]),call_result[1])
+    # if type(call_result) is tuple:
+    #     for i in range(len(call_result)):
+    #         call_result[i] = HWC3(call_result[i])
+    # else:
+    #     call_result = HWC3(call_result)
+    return HWC3(call_result)
+
 
 class CannyEdgePreprocesor:
     @classmethod
diff --git a/custom_nodes/annotator/uniformer/__init__.py b/custom_nodes/annotator/uniformer/__init__.py
index 6be429542..bd20c5722 100644
--- a/custom_nodes/annotator/uniformer/__init__.py
+++ b/custom_nodes/annotator/uniformer/__init__.py
@@ -14,7 +14,7 @@ class UniformerDetector:
         if not os.path.exists(modelpath):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(checkpoint_file, model_dir=annotator_ckpts_path)
-        config_file = os.path.join(os.path.dirname(annotator_ckpts_path), "uniformer", "exp", "upernet_global_small", "config.py")
+        config_file = os.path.join(os.path.dirname(annotator_ckpts_path).replace("models", "custom_nodes/annotator"), "uniformer", "exp", "upernet_global_small", "config.py")
         self.model = init_segmentor(config_file, modelpath).cuda()
 
     def __call__(self, img):