diff --git a/__init__.py b/__init__.py index 0daf8ac..730a1fa 100644 --- a/__init__.py +++ b/__init__.py @@ -52,6 +52,7 @@ NODE_CLASS_MAPPINGS = { "PM_FaceShapMatch": FaceShapMatchPM, "PM_SuperColorTransfer": SuperColorTransferPM, "PM_SuperMakeUpTransfer": SuperMakeUpTransferPM, + "PM_Similarity": SimilarityPM, } NODE_DISPLAY_NAME_MAPPINGS = { "PM_RetinaFace": "RetinaFace PM", @@ -73,6 +74,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PM_FaceShapMatch": "FaceShapMatch PM", "PM_SuperColorTransfer": "SuperColorTransfer PM", "PM_SuperMakeUpTransfer": "SuperMakeUpTransfer PM", + "PM_Similarity": "Similarity PM", + } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/portrait/model_holder.py b/portrait/model_holder.py index b10374b..2f74f06 100644 --- a/portrait/model_holder.py +++ b/portrait/model_holder.py @@ -16,6 +16,7 @@ skin_retouching = None portrait_enhancement = None psgan_interface = None real_gan_sr = None +face_recognition = None def get_retinaface_detection(): global retinaface_detection @@ -72,3 +73,10 @@ def get_pagan_interface(): makeup_transfer_model_path = os.path.join(models_path, "makeup_transfer.pth") psgan_interface = PSGAN_Inference("cuda", makeup_transfer_model_path, get_retinaface_detection(), get_face_skin(), face_landmarks_model_path) return psgan_interface + +def get_face_recognition(): + global face_recognition + if face_recognition is None: + face_recognition = pipeline("face_recognition", model="bubbliiiing/cv_retinafce_recognition", model_revision="v1.0.3") + return face_recognition + diff --git a/portrait/nodes.py b/portrait/nodes.py index b004fb5..4cb475b 100644 --- a/portrait/nodes.py +++ b/portrait/nodes.py @@ -303,7 +303,6 @@ class ImageScaleShortPM: CATEGORY = "protrait/other" - def image_scale_short(self, image, size, crop_face): input_image = tensor_to_img(image) short_side = min(input_image.width, input_image.height) @@ -485,3 +484,31 @@ class SuperMakeUpTransferPM: box_size_transfer = transfer_image.resize([box_width, box_height], Image.Resampling.LANCZOS) origin_np[avatar_box[1]:avatar_box[3], avatar_box[0]:avatar_box[2], :] = img_to_np(box_size_transfer) return (np_to_tensor(origin_np),) + +class SimilarityPM: + @classmethod + def INPUT_TYPES(s): + return \ + { + "required": { + "main_image": ("IMAGE",), + "compare_image": ("IMAGE",), + "model": (["sim"],), + }, + "optional": { + "avatar_box": ("BOX",), + }, + } + + RETURN_TYPES = ("FLOAT",) + + FUNCTION = "similarity_compare" + CATEGORY = "protrait/model" + + def similarity_compare(self, main_image, compare_image, model): + score = None + if model == "sim": + root_embedding = get_face_recognition(dict(user=Image.fromarray(np.uint8(main_image))))[OutputKeys.IMG_EMBEDDING] + compare_embedding = face_recognition(dict(user=Image.fromarray(np.uint8(compare_image))))[OutputKeys.IMG_EMBEDDING] + score = float(np.dot(root_embedding, np.transpose(compare_embedding))[0][0]) + return (score,)