Fix Openpose and Segmentation

Fix Openpose: Add code explicitly for openpose
Segmentations (Uniformer): Fix path mismatch (config is not in models folder)
This commit is contained in:
NiggoJaecha 2023-02-22 23:04:19 +01:00
parent 4197f7d7c2
commit 9075ec965b
2 changed files with 10 additions and 7 deletions

View File

@ -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

View File

@ -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):