retain face choose face

README.md

@@ -19,6 +19,7 @@ If you have any questions or suggestions, you can reach us through:
 3. Add PM_MakeUpTransfer node. same as easyphoto makeup transfer.
 4. Add a super-resolution model to the PM_PortraitEnhancement node. This super-resolution model can not highlight faces.
 5. Add v1.1.0 workflow
+6. RetinaFace supports face selection
 
 ## V1.0.0 Update
 1. Added log for model downloads.
@@ -68,6 +69,7 @@ Click "Load" in the right panel of ComfyUI and select the ./workflow/easyphoto_w
 * RetainFace PM: Perform matting using models from Model Scope. [Link](https://www.modelscope.cn/models/damo/cv_resnet50_face-detection_retinaface/summary)
   * image: Input image
   * multi_user_facecrop_ratio: Multiplicative factor for extracting the head region.
+  * face_index : Choose which face
 
 * FaceFusion PM: Merge faces from two images.
   * image: Input image

README_zh-CN.md

@@ -21,6 +21,7 @@ English | [简体中文](./README_zh-CN.md)
 3. 增加 PM_MakeUpTransfer节点 与easyphoto的MakeupTransfer一致
 4. PM_PortraitEnhancement节点增加一种超分模型，此超分模型可以对人脸不做高光
 5. 增加v1.1.0 workflow
+6. RetinaFace 支持选择人脸
 
 ## v1.0.0 更新
 
@@ -71,6 +72,7 @@ Easyphoto工作位置: [./workflow/easyphoto.json](./workflows/easyphoto.json )
 * RetainFace PM:使用Model Scope中的模型进行抠图 [链接](https://www.modelscope.cn/models/damo/cv_resnet50_face-detection_retinaface/summary)
 	* image：输入图像
 	* multi_user_facecrop_ratio：提取头像区域的倍数
+    * face_index : 选择第几个人脸
 * FaceFusion PM:将两张图像的人脸进行融合
 	* image：输入图像
 	* user_image：要融合的头像

portrait/nodes.py

@@ -13,7 +13,8 @@ class RetinaFacePM:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {"image": ("IMAGE",),
-                             "multi_user_facecrop_ratio": ("FLOAT", {"default": 1, "min": 0, "max": 10, "step": 0.01})
+                             "multi_user_facecrop_ratio": ("FLOAT", {"default": 1, "min": 0, "max": 10, "step": 0.01}),
+                             "face_index": ("INT", {"default": 0, "min": 0, "max": 10, "step": 1})
                              }}
 
     RETURN_TYPES = ("IMAGE", "MASK", "BOX")
@@ -21,12 +22,14 @@ class RetinaFacePM:
     FUNCTION = "retain_face"
     CATEGORY = "protrait/model"
 
-    def retain_face(self, image, multi_user_facecrop_ratio):
+    def retain_face(self, image, multi_user_facecrop_ratio, face_index):
         np_image = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
         image = Image.fromarray(np_image)
-        retinaface_boxes, retinaface_keypoints, retinaface_masks, retinaface_tensor = call_face_crop(get_retinaface_detection(), image, multi_user_facecrop_ratio)
-        crop_image = image.crop(retinaface_boxes[0])
-        return (img_to_tensor(crop_image), retinaface_tensor, retinaface_boxes[0])
+        retinaface_boxes, retinaface_keypoints, retinaface_masks, retinaface_mask_nps = call_face_crop(get_retinaface_detection(), image, multi_user_facecrop_ratio)
+        crop_image = image.crop(retinaface_boxes[face_index])
+        retinaface_mask = np_to_mask(retinaface_mask_nps[face_index])
+        retinaface_boxe = retinaface_boxes[face_index]
+        return (img_to_tensor(crop_image), retinaface_mask, retinaface_boxe)
 
 class FaceFusionPM:
 

portrait/utils/face_process_utils.py

@@ -65,9 +65,8 @@ def safe_get_box_mask_keypoints(image, retinaface_result, crop_ratio, face_seg,
         retinaface_boxs = [retinaface_boxs[index] for index in argindex]
         retinaface_keypoints = [retinaface_keypoints[index] for index in argindex]
         retinaface_mask_pils = [retinaface_mask_pils[index] for index in argindex]
-        retinaface_mask_np = [retinaface_masks[index] for index in argindex]
-        mask_tensor = np_to_mask(retinaface_mask_np[0])
-        return retinaface_boxs, retinaface_keypoints, retinaface_mask_pils, mask_tensor
+        retinaface_mask_nps = [retinaface_masks[index] for index in argindex]
+        return retinaface_boxs, retinaface_keypoints, retinaface_mask_pils, retinaface_mask_nps
 
     else:
         retinaface_box = np.array([])
@@ -120,9 +119,9 @@ def call_face_crop(retinaface_detection, image, crop_ratio, prefix="tmp"):
     # retinaface detect 
     retinaface_result = retinaface_detection(image)
     # get mask and keypoints
-    retinaface_box, retinaface_keypoints, retinaface_mask_pil, retinaface_mask_tensor = safe_get_box_mask_keypoints(image, retinaface_result, crop_ratio, None, "crop")
+    retinaface_box, retinaface_keypoints, retinaface_mask_pil, retinaface_mask_nps = safe_get_box_mask_keypoints(image, retinaface_result, crop_ratio, None, "crop")
 
-    return retinaface_box, retinaface_keypoints, retinaface_mask_pil, retinaface_mask_tensor
+    return retinaface_box, retinaface_keypoints, retinaface_mask_pil, retinaface_mask_nps
 
 def color_transfer(sc, dc):
     """