From a54881ac60e5e5316a2f03efd486f7cf4e1247f6 Mon Sep 17 00:00:00 2001 From: tototianhao <5173315+toto000888@user.noreply.gitee.com> Date: Sat, 6 Jan 2024 15:27:27 +0800 Subject: [PATCH] add face swap and seg node --- __init__.py | 6 +- facechain/common/model_processor.py | 76 +++++++++++++ facechain/model_holder.py | 3 +- facechain/nodes.py | 161 ++-------------------------- 4 files changed, 91 insertions(+), 155 deletions(-) diff --git a/__init__.py b/__init__.py index 4be0c79..5b12248 100644 --- a/__init__.py +++ b/__init__.py @@ -11,20 +11,18 @@ from .facechain.style_loader_node import * NODE_CLASS_MAPPINGS = { "FC FaceFusion": FCFaceFusion, - "FC StyleLoraLoad": FCStyleLoraLoad, "FC FaceDetectCrop": FaceDetectCrop, "FC FaceSegment": FCFaceSegment, - "FC ReplaceImage": FCReplaceImage, + "FC FaceSegAndReplace": FCFaceSegAndReplace, "FC CropBottom": FCCropBottom, "FC CropAndPaste": FCCropAndPaste, "FC MaskOP": FCMaskOP, } NODE_DISPLAY_NAME_MAPPINGS = { "FC FaceFusion": "FC FaceFusion", - "FC StyleLoraLoad": "FC StyleLoraLoad", "FC FaceDetectCrop": "FC FaceDetectCrop", "FC FaceSegment": "FC FaceSegment", - "FC ReplaceImage": "FC ReplaceImage", + "FC FaceSegAndReplace": "FC FaceSegAndReplace", "FC CropBottom": "FC CropBottom", "FC CropAndPaste": "FC CropAndPaste", "FC MaskOP": "FC MaskOP", diff --git a/facechain/common/model_processor.py b/facechain/common/model_processor.py index dca8abf..7f3d73a 100644 --- a/facechain/common/model_processor.py +++ b/facechain/common/model_processor.py @@ -1,5 +1,7 @@ import cv2 import numpy as np +from modelscope.outputs import OutputKeys + from facechain.model_holder import * from facechain.utils.convert_utils import * @@ -64,3 +66,77 @@ def facechain_detect_crop(source_image_pil, face_index, crop_ratio, mode): return inpaint_img, mask, bbox, points_array, else: raise RuntimeError('模式错误') + + +def segment(img, ksize=0, eyeh=0, ksize1=0, include_neck=False, warp_mask=None, return_human=False): + result = get_segmentation()(img) + masks = result['masks'] + scores = result['scores'] + labels = result['labels'] + if len(masks) == 0: + return + h, w = masks[0].shape + mask_face = np.zeros((h, w)) + mask_hair = np.zeros((h, w)) + mask_neck = np.zeros((h, w)) + mask_cloth = np.zeros((h, w)) + mask_human = np.zeros((h, w)) + for i in range(len(labels)): + if scores[i] > 0.8: + if labels[i] == 'Torso-skin': + mask_neck += masks[i] + elif labels[i] == 'Face': + mask_face += masks[i] + elif labels[i] == 'Human': + mask_human += masks[i] + elif labels[i] == 'Hair': + mask_hair += masks[i] + elif labels[i] == 'UpperClothes' or labels[i] == 'Coat': + mask_cloth += masks[i] + mask_face = np.clip(mask_face, 0, 1) + mask_hair = np.clip(mask_hair, 0, 1) + mask_neck = np.clip(mask_neck, 0, 1) + mask_cloth = np.clip(mask_cloth, 0, 1) + mask_human = np.clip(mask_human, 0, 1) + soft_mask = 0 + if np.sum(mask_face) > 0: + soft_mask = np.clip(mask_face, 0, 1) + if ksize1 > 0: + kernel_size1 = int(np.sqrt(np.sum(soft_mask)) * ksize1) + kernel1 = np.ones((kernel_size1, kernel_size1)) + soft_mask = cv2.dilate(soft_mask, kernel1, iterations=1) + if ksize > 0: + kernel_size = int(np.sqrt(np.sum(soft_mask)) * ksize) + kernel = np.ones((kernel_size, kernel_size)) + soft_mask_dilate = cv2.dilate(soft_mask, kernel, iterations=1) + if warp_mask is not None: + soft_mask_dilate = soft_mask_dilate * (np.clip(soft_mask + warp_mask[:, :, 0], 0, 1)) + if eyeh > 0: + soft_mask = np.concatenate((soft_mask[:eyeh], soft_mask_dilate[eyeh:]), axis=0) + else: + soft_mask = soft_mask_dilate + else: + if ksize1 > 0: + kernel_size1 = int(np.sqrt(np.sum(soft_mask)) * ksize1) + kernel1 = np.ones((kernel_size1, kernel_size1)) + soft_mask = cv2.dilate(mask_face, kernel1, iterations=1) + else: + soft_mask = mask_face + if include_neck: + soft_mask = np.clip(soft_mask + mask_neck, 0, 1) + + if return_human: + mask_human = cv2.GaussianBlur(mask_human, (21, 21), 0) * mask_human + return soft_mask, mask_human + else: + return soft_mask + + +def face_fusing_seg_replace(image, template_face): + image_face_fusion = pipeline('face_fusion_torch', model='damo/cv_unet_face_fusion_torch', model_revision='v1.0.5') + result = image_face_fusion(dict(template=image, user=template_face))[OutputKeys.OUTPUT_IMG] + debug(result) + face_mask = segment(image, ksize=0.1) + result = (result * face_mask[:, :, None] + np.array(image)[:, :, ::-1] * (1 - face_mask[:, :, None])).astype(np.uint8) + debug(result) + return result diff --git a/facechain/model_holder.py b/facechain/model_holder.py index 6e05990..fd8c1d3 100644 --- a/facechain/model_holder.py +++ b/facechain/model_holder.py @@ -23,11 +23,12 @@ def get_face_detection(): face_detection = pipeline(task=Tasks.face_detection, model='damo/cv_ddsar_face-detection_iclr23-damofd', model_revision='v1.1') return face_detection + image_face_fusion = pipeline('face_fusion_torch', model='damo/cv_unet_face_fusion_torch', model_revision='v1.0.5') def get_image_face_fusion(): global image_face_fusion if image_face_fusion is None: - image_face_fusion = pipeline(Tasks.image_face_fusion, model='damo/cv_unet-image-face-fusion_damo', model_revision='v1.3') + image_face_fusion = pipeline('face_fusion_torch', model='damo/cv_unet_face_fusion_torch', model_revision='v1.0.5') return image_face_fusion diff --git a/facechain/nodes.py b/facechain/nodes.py index 2f454f9..50493f7 100644 --- a/facechain/nodes.py +++ b/facechain/nodes.py @@ -9,7 +9,7 @@ from modelscope.outputs import OutputKeys from facechain.model_holder import * from facechain.utils.img_utils import * from facechain.utils.convert_utils import * -from facechain.common.model_processor import facechain_detect_crop +from facechain.common.model_processor import * class FCFaceFusion: @@ -31,7 +31,7 @@ class FCFaceFusion: fusion_image = tensor_to_img(fusion_image) result_image = get_image_face_fusion()(dict(template=source_image, user=fusion_image))[OutputKeys.OUTPUT_IMG] result_image = Image.fromarray(cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)) - return (img_to_tensor(result_image),) + return (image_to_tensor(result_image),) class FaceDetectCrop: @@ -59,58 +59,6 @@ class FaceDetectCrop: return (image_to_tensor(corp_img_pil), mask_np3_to_mask_tensor(mask), bbox, points_array,) -class FCCropMask: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "face_box": ("BOX",) - } - } - - RETURN_TYPES = ("IMAGE", "MASK",) - FUNCTION = "crop_mask" - CATEGORY = "facechain/mask" - - def crop_mask(self, image, face_box): - image = tensor_to_np(image) - inpaint_img_large = image - mask_large = np.ones_like(inpaint_img_large) - mask_large1 = np.zeros_like(inpaint_img_large) - h, w, _ = inpaint_img_large.shape - face_ratio = 0.45 - cropl = int(max(face_box[3] - face_box[1], face_box[2] - face_box[0]) / face_ratio / 2) - cx = int((face_box[2] + face_box[0]) / 2) - cy = int((face_box[1] + face_box[3]) / 2) - cropup = min(cy, cropl) - cropbo = min(h - cy, cropl) - crople = min(cx, cropl) - cropri = min(w - cx, cropl) - inpaint_img = np.pad(inpaint_img_large[cy - cropup:cy + cropbo, cx - crople:cx + cropri], ((cropl - cropup, cropl - cropbo), (cropl - crople, cropl - cropri), (0, 0)), - 'constant') - inpaint_img = cv2.resize(inpaint_img, (512, 512)) - inpaint_img = Image.fromarray(cv2.cvtColor(inpaint_img[:, :, ::-1], cv2.COLOR_BGR2RGB)) - mask_large1[cy - cropup:cy + cropbo, cx - crople:cx + cropri] = 1 - mask_large = mask_large * mask_large1 - return (image_to_tensor(inpaint_img), mask_np3_to_mask_tensor(mask_large)) - - -class FCFaceSwap(): - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "face_box": ("BOX",) - } - } - - RETURN_TYPES = ("IMAGE", "MASK",) - FUNCTION = "crop_mask" - CATEGORY = "facechain/mask" - - class FCFaceSegment: @classmethod def INPUT_TYPES(s): @@ -124,78 +72,14 @@ class FCFaceSegment: FUNCTION = "fc_segment" CATEGORY = "facechain/model" - def segment(self, segmentation_pipeline, img, ksize=0, eyeh=0, ksize1=0, include_neck=False, warp_mask=None, return_human=False): - if True: - result = segmentation_pipeline(img) - masks = result['masks'] - scores = result['scores'] - labels = result['labels'] - if len(masks) == 0: - return - h, w = masks[0].shape - mask_face = np.zeros((h, w)) - mask_hair = np.zeros((h, w)) - mask_neck = np.zeros((h, w)) - mask_cloth = np.zeros((h, w)) - mask_human = np.zeros((h, w)) - for i in range(len(labels)): - if scores[i] > 0.8: - if labels[i] == 'Torso-skin': - mask_neck += masks[i] - elif labels[i] == 'Face': - mask_face += masks[i] - elif labels[i] == 'Human': - mask_human += masks[i] - elif labels[i] == 'Hair': - mask_hair += masks[i] - elif labels[i] == 'UpperClothes' or labels[i] == 'Coat': - mask_cloth += masks[i] - mask_face = np.clip(mask_face, 0, 1) - mask_hair = np.clip(mask_hair, 0, 1) - mask_neck = np.clip(mask_neck, 0, 1) - mask_cloth = np.clip(mask_cloth, 0, 1) - mask_human = np.clip(mask_human, 0, 1) - soft_mask = 0 - if np.sum(mask_face) > 0: - soft_mask = np.clip(mask_face, 0, 1) - if ksize1 > 0: - kernel_size1 = int(np.sqrt(np.sum(soft_mask)) * ksize1) - kernel1 = np.ones((kernel_size1, kernel_size1)) - soft_mask = cv2.dilate(soft_mask, kernel1, iterations=1) - if ksize > 0: - kernel_size = int(np.sqrt(np.sum(soft_mask)) * ksize) - kernel = np.ones((kernel_size, kernel_size)) - soft_mask_dilate = cv2.dilate(soft_mask, kernel, iterations=1) - if warp_mask is not None: - soft_mask_dilate = soft_mask_dilate * (np.clip(soft_mask + warp_mask[:, :, 0], 0, 1)) - if eyeh > 0: - soft_mask = np.concatenate((soft_mask[:eyeh], soft_mask_dilate[eyeh:]), axis=0) - else: - soft_mask = soft_mask_dilate - else: - if ksize1 > 0: - kernel_size1 = int(np.sqrt(np.sum(soft_mask)) * ksize1) - kernel1 = np.ones((kernel_size1, kernel_size1)) - soft_mask = cv2.dilate(mask_face, kernel1, iterations=1) - else: - soft_mask = mask_face - if include_neck: - soft_mask = np.clip(soft_mask + mask_neck, 0, 1) - - if return_human: - mask_human = cv2.GaussianBlur(mask_human, (21, 21), 0) * mask_human - return soft_mask, mask_human - else: - return soft_mask - def fc_segment(self, source_image): pil_source_image = tensor_to_img(source_image) - mask = self.segment(get_segmentation(), pil_source_image, ksize=0.1) + mask = segment(pil_source_image, ksize=0.1) seg_image = tensor_to_np(source_image) * mask[:, :, None] return (image_np_to_image_tensor(seg_image), mask_np2_to_mask_tensor(mask),) -class FCReplaceImage: +class FCFaceSegAndReplace: @classmethod def INPUT_TYPES(s): return { @@ -208,37 +92,14 @@ class FCReplaceImage: } RETURN_TYPES = ("IMAGE",) - FUNCTION = "replace_image" + FUNCTION = "face_swap" CATEGORY = "facechain/model" - def replace_image(self, source_image, replace_image, face_box, mask): - face_ratio = 0.45 - h, w, _ = replace_image.shape - cropl = int(max(face_box[3] - face_box[1], face_box[2] - face_box[0]) / face_ratio / 2) - cx = int((face_box[2] + face_box[0]) / 2) - cy = int((face_box[1] + face_box[3]) / 2) - cropup = min(cy, cropl) - cropbo = min(h - cy, cropl) - crople = min(cx, cropl) - cropri = min(w - cx, cropl) - ksize = int(10 * cropl / 256) - rst_gen = cv2.resize(replace_image, (cropl * 2, cropl * 2)) - rst_crop = rst_gen[cropl - cropup:cropl + cropbo, cropl - crople:cropl + cropri] - print(rst_crop.shape) - inpaint_img_rst = np.zeros_like(source_image) - print('Start pasting.') - inpaint_img_rst[cy - cropup:cy + cropbo, cx - crople:cx + cropri] = rst_crop - print('Fininsh pasting.') - print(inpaint_img_rst.shape, mask.shape, source_image.shape) - mask_large = mask.astype(np.float32) - kernel = np.ones((ksize * 2, ksize * 2)) - mask_large1 = cv2.erode(mask_large, kernel, iterations=1) - mask_large1 = cv2.GaussianBlur(mask_large1, (int(ksize * 1.8) * 2 + 1, int(ksize * 1.8) * 2 + 1), 0) - mask_large1[face_box[1]:face_box[3], face_box[0]:face_box[2]] = 1 - mask_large = mask_large * mask_large1 - final_inpaint_rst = (inpaint_img_rst.astype(np.float32) * mask_large.astype(np.float32) + source_image.astype(np.float32) * (1.0 - mask_large.astype(np.float32))).astype( - np.uint8) - return (final_inpaint_rst,) + def face_swap(self, source_image, replace_image): + pil_source_image = image_to_tensor(source_image) + pil_replace_image = image_to_tensor(replace_image) + image = face_fusing_seg_replace(pil_source_image, pil_replace_image) + return (image_np_to_image_tensor(image),) class FCCropBottom: @@ -258,7 +119,7 @@ class FCCropBottom: def crop_bottom(self, source_image, width): source_image = tensor_to_img(source_image) crop_result = crop_bottom(source_image, width) - return (img_to_tensor(crop_result),) + return (image_to_tensor(crop_result),) class FCCropAndPaste: