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): def common_annotator_call(annotator_callback, tensor_image, *args):
call_result = annotator_callback(img_tensor_to_np(tensor_image), *args) call_result = annotator_callback(img_tensor_to_np(tensor_image), *args)
if type(call_result) is tuple: if type(annotator_callback) is openpose.OpenposeDetector:
for i in range(len(call_result)): return (HWC3(call_result[0]),call_result[1])
call_result[i] = HWC3(call_result[i]) # if type(call_result) is tuple:
else: # for i in range(len(call_result)):
call_result = HWC3(call_result) # call_result[i] = HWC3(call_result[i])
return call_result # else:
# call_result = HWC3(call_result)
return HWC3(call_result)
class CannyEdgePreprocesor: class CannyEdgePreprocesor:
@classmethod @classmethod

View File

@ -14,7 +14,7 @@ class UniformerDetector:
if not os.path.exists(modelpath): if not os.path.exists(modelpath):
from basicsr.utils.download_util import load_file_from_url from basicsr.utils.download_util import load_file_from_url
load_file_from_url(checkpoint_file, model_dir=annotator_ckpts_path) 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() self.model = init_segmentor(config_file, modelpath).cuda()
def __call__(self, img): def __call__(self, img):