update sim

__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']

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
+

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,)