From 4abaf4cd92a5c4e5799c63200bd0080db99c19ee Mon Sep 17 00:00:00 2001 From: toto Date: Wed, 25 Oct 2023 11:50:21 +0800 Subject: [PATCH] retain face choose face --- README.md | 2 ++ README_zh-CN.md | 2 ++ portrait/nodes.py | 13 ++++++++----- portrait/utils/face_process_utils.py | 9 ++++----- 4 files changed, 16 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 0813dab..fbe4579 100644 --- a/README.md +++ b/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 diff --git a/README_zh-CN.md b/README_zh-CN.md index c855abf..8709170 100644 --- a/README_zh-CN.md +++ b/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：要融合的头像 diff --git a/portrait/nodes.py b/portrait/nodes.py index bbb97de..4a1c92a 100644 --- a/portrait/nodes.py +++ b/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: diff --git a/portrait/utils/face_process_utils.py b/portrait/utils/face_process_utils.py index 3380891..7a5458b 100644 --- a/portrait/utils/face_process_utils.py +++ b/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): """