diff --git a/__init__.py b/__init__.py index 5b12248..a665abe 100644 --- a/__init__.py +++ b/__init__.py @@ -6,36 +6,42 @@ import threading root_path = os.path.dirname(__file__) parent_dir = os.path.dirname(root_path) sys.path.append(root_path) -from .facechain.nodes import * -from .facechain.style_loader_node import * +from facechain.nodes import * +from facechain.style_loader_node import * NODE_CLASS_MAPPINGS = { - "FC FaceFusion": FCFaceFusion, "FC FaceDetectCrop": FaceDetectCrop, + "FC FaceFusion": FCFaceFusion, "FC FaceSegment": FCFaceSegment, - "FC FaceSegAndReplace": FCFaceSegAndReplace, + "FC FaceSegAndReplace": FCFaceFusionAndSegReplace, + "FC RemoveCannyFace": FCRemoveCannyFace, "FC CropBottom": FCCropBottom, + "FC ReplaceByMask": FCReplaceByMask, "FC CropAndPaste": FCCropAndPaste, "FC MaskOP": FCMaskOP, } NODE_DISPLAY_NAME_MAPPINGS = { - "FC FaceFusion": "FC FaceFusion", "FC FaceDetectCrop": "FC FaceDetectCrop", + "FC FaceFusion": "FC FaceFusion", "FC FaceSegment": "FC FaceSegment", "FC FaceSegAndReplace": "FC FaceSegAndReplace", + "FC RemoveCannyFace": "FC RemoveCannyFace", "FC CropBottom": "FC CropBottom", + "FC ReplaceByMask": "FC ReplaceByMask", "FC CropAndPaste": "FC CropAndPaste", "FC MaskOP": "FC MaskOP", } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + # install_model = ["python-slugify==8.0.1", "modelscope", "controlnet_aux==0.0.6", "onnxruntime==1.15.1", "mmcv==1.7.0", "mmdet==2.26.0", "mediapipe==0.10.3", "edge_tts"] def handle_stream(stream, prefix): for line in stream: print(prefix, line, end="") + def run_script(cmd, cwd='.'): process = subprocess.Popen(cmd, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1) stdout_thread = threading.Thread(target=handle_stream, args=(process.stdout, "")) @@ -46,6 +52,7 @@ def run_script(cmd, cwd='.'): stderr_thread.join() return process.wait() + if os.path.basename(parent_dir) == "custom_nodes": print("## installing facechain dependencies") requirements_path = os.path.join(root_path, "requirements.txt") diff --git a/facechain/common/model_processor.py b/facechain/common/model_processor.py index 7f3d73a..a80b409 100644 --- a/facechain/common/model_processor.py +++ b/facechain/common/model_processor.py @@ -1,6 +1,7 @@ import cv2 import numpy as np from modelscope.outputs import OutputKeys +from skimage import transform from facechain.model_holder import * from facechain.utils.convert_utils import * @@ -10,7 +11,7 @@ def debug(*args): print(f"==== face chain debug ====", *args) -def facechain_detect_crop(source_image_pil, face_index, crop_ratio, mode): +def facechain_detect_crop(source_image_pil, face_index=0, crop_ratio=1, mode='normal'): det_result = get_face_detection()(source_image_pil) mask = np.zeros_like(source_image_pil) bboxes = det_result['boxes'] @@ -37,7 +38,7 @@ def facechain_detect_crop(source_image_pil, face_index, crop_ratio, mode): points_array[k, 0] = keypoint[2 * k] points_array[k, 1] = keypoint[2 * k + 1] debug(f'mode = ', mode) - if mode == "real seg": + if mode == "normal": bbox[0] = np.clip(np.array(bbox[0], np.int32) - face_w * (crop_ratio - 1) / 2, 0, w - 1) bbox[1] = np.clip(np.array(bbox[1], np.int32) - face_h * (crop_ratio - 1) / 2, 0, h - 1) bbox[2] = np.clip(np.array(bbox[2], np.int32) + face_w * (crop_ratio - 1) / 2, 0, w - 1) @@ -68,11 +69,11 @@ def facechain_detect_crop(source_image_pil, face_index, crop_ratio, mode): 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'] +def segment(img, ksize=0, eyeh=0, ksize1=0, include_neck=False, warp_mask=None, return_human=True): + seg_image = get_segmentation()(img) + masks = seg_image['masks'] + scores = seg_image['scores'] + labels = seg_image['labels'] if len(masks) == 0: return h, w = masks[0].shape @@ -124,19 +125,55 @@ def segment(img, ksize=0, eyeh=0, ksize1=0, include_neck=False, warp_mask=None, soft_mask = mask_face if include_neck: soft_mask = np.clip(soft_mask + mask_neck, 0, 1) - + np_image = image_to_np(img) + img.crop() + seg_image_np = np_image * soft_mask[:, :, None] + seg_image = Image.fromarray(seg_image_np.astype(np.uint8)) if return_human: mask_human = cv2.GaussianBlur(mask_human, (21, 21), 0) * mask_human - return soft_mask, mask_human + return seg_image, soft_mask, mask_human else: - return soft_mask + # 返回一个是PIL，一个是np的二维数组 + return seg_image, 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 +def face_fusion(image, fusion_image): + cv_fusion_result = get_image_face_fusion()(dict(template=image, user=fusion_image))[OutputKeys.OUTPUT_IMG] + np_fusion_result = cv2.cvtColor(cv_fusion_result, cv2.COLOR_BGR2RGB) + debug("cv_fusion_result", cv_fusion_result.shape) + return np_fusion_result + + +def face_fusing_seg_replace(image, replace_image): + np_fusion_result = face_fusion(image, replace_image) + _, face_mask, _ = segment(image, ksize=0.1) + cv_replace_result = (np_fusion_result * face_mask[:, :, None] + np.array(image) * (1 - face_mask[:, :, None])).astype(np.uint8) + debug("cv_replace_result", cv_replace_result.shape) + return np_fusion_result, cv_replace_result + + +def crop_and_paste(Source_image, Source_image_mask, Target_image, Source_Five_Point, Target_Five_Point, Source_box, use_warp=True): + debug(f"crop and paste", Source_image, Source_image_mask, Target_image, Source_Five_Point, Target_Five_Point, Source_box) + if use_warp: + Source_Five_Point = np.reshape(Source_Five_Point, [5, 2]) - np.array(Source_box[:2]) + Target_Five_Point = np.reshape(Target_Five_Point, [5, 2]) + + Crop_Source_image = Source_image.crop(np.int32(Source_box)) + Crop_Source_image_mask = Source_image_mask.crop(np.int32(Source_box)) + Source_Five_Point, Target_Five_Point = np.array(Source_Five_Point), np.array(Target_Five_Point) + + tform = transform.SimilarityTransform() + tform.estimate(Source_Five_Point, Target_Five_Point) + M = tform.params[0:2, :] + + warped = cv2.warpAffine(np.array(Crop_Source_image), M, np.shape(Target_image)[:2][::-1], borderValue=0.0) + warped_mask = cv2.warpAffine(np.array(Crop_Source_image_mask), M, np.shape(Target_image)[:2][::-1], borderValue=0.0) + + mask = np.float32(warped_mask == 0) + debug('target shape', np.float32(Target_image).shape) + debug('warped shape', np.float32(warped).shape) + output = mask * np.float32(Target_image) + (1 - mask) * np.float32(warped) + else: + mask = np.float32(np.array(Source_image_mask) == 0) + output = mask * np.float32(Target_image) + (1 - mask) * np.float32(Source_image) + return output, mask diff --git a/facechain/model_holder.py b/facechain/model_holder.py index fd8c1d3..7e50e57 100644 --- a/facechain/model_holder.py +++ b/facechain/model_holder.py @@ -23,7 +23,6 @@ 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 diff --git a/facechain/nodes.py b/facechain/nodes.py index 50493f7..f32efb9 100644 --- a/facechain/nodes.py +++ b/facechain/nodes.py @@ -29,9 +29,8 @@ class FCFaceFusion: def image_face_fusion(self, source_image, fusion_image): source_image = tensor_to_img(source_image) 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 (image_to_tensor(result_image),) + result_image = face_fusion(source_image, fusion_image) + return (image_np_to_image_tensor(result_image),) class FaceDetectCrop: @@ -45,7 +44,7 @@ class FaceDetectCrop: "source_image": ("IMAGE",), "face_index": ("INT", {"default": 0, "min": 0, "max": 10, "step": 1}), "crop_ratio": ("FLOAT", {"default": 1.0, "min": 0, "max": 10, "step": 0.1}), - "mode": (["real seg", "square 512 width heigh"],), + "mode": (["normal", "square 512 width heigh"],), } } @@ -65,41 +64,71 @@ class FCFaceSegment: return { "required": { "source_image": ("IMAGE",), - } + + }, + "optional": { + "ksize": ("FLOAT", {"default": 0, "min": 0, "max": 10, "step": 0.1}), + "ksize1": ("FLOAT", {"default": 0, "min": 0, "max": 10, "step": 0.1}), + "include_neck": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + "warp_mask": ("MASK",), + }, } - RETURN_TYPES = ("IMAGE", "MASK",) + RETURN_TYPES = ("IMAGE", "MASK", "MASK") + RETURN_NAMES = ("seg_image", "soft_mask", "human_mask") FUNCTION = "fc_segment" CATEGORY = "facechain/model" - def fc_segment(self, source_image): + def fc_segment(self, source_image, ksize=0, ksize1=0, include_neck=False, warp_mask=None, ): pil_source_image = tensor_to_img(source_image) - 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),) + seg_image, mask, human_mask = segment(pil_source_image, ksize, ksize1, include_neck, warp_mask, True) + return image_to_tensor(seg_image), mask_np2_to_mask_tensor(mask), mask_np2_to_mask_tensor(human_mask) -class FCFaceSegAndReplace: +class FCFaceFusionAndSegReplace: @classmethod def INPUT_TYPES(s): return { "required": { "source_image": ("IMAGE",), "replace_image": ("IMAGE",), - "face_box": ("BOX",), - "mask": ("MASK",) } } - RETURN_TYPES = ("IMAGE",) + RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_NAMES = ("fusion", "fusion seg replace") FUNCTION = "face_swap" CATEGORY = "facechain/model" 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),) + pil_source_image = tensor_to_img(source_image) + pil_replace_image = tensor_to_img(replace_image) + cv_fusion_result, cv_replace_result = face_fusing_seg_replace(pil_source_image, pil_replace_image) + return (image_np_to_image_tensor(cv_fusion_result), image_np_to_image_tensor(cv_replace_result),) + + +class FCRemoveCannyFace: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source_image": ("IMAGE",), + "canny_image": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "remove_canny_face" + CATEGORY = "facechain/model" + + def remove_canny_face(self, source_image, canny_image): + pil_source_image = tensor_to_img(source_image) + np_canny_image = tensor_to_np(canny_image) + corp_img_pil, _, _, points_array = facechain_detect_crop(pil_source_image, 0, 1.1, 'normal') + eye_height = int((points_array[0, 1] + points_array[1, 1]) / 2) + _, mask, _ = segment(pil_source_image, ksize=0.05, eyeh=eye_height) + canny_image = (np_canny_image * (1.0 - mask[:, :, None])).astype(np.uint8) + return (image_np_to_image_tensor(canny_image),) class FCCropBottom: @@ -108,7 +137,7 @@ class FCCropBottom: return { "required": { "source_image": ("IMAGE",), - "face_index": ("INT", {"default": 0, "min": 0, "max": 10, "step": 1}) + "width": ("INT", {"default": 512, "min": 0, "max": 2048, "step": 1}) } } @@ -122,49 +151,80 @@ class FCCropBottom: return (image_to_tensor(crop_result),) -class FCCropAndPaste: +class FCEdgeAdd: @classmethod def INPUT_TYPES(s): return { "required": { "source_image": ("IMAGE",), - "source_image_mask": ("MASK",), - "source_box": ("BOX",), - "source_five_point": ("KEY_POINT",), - "target_image": ("IMAGE",), - "target_five_point": ("KEY_POINT",), + "edge_add_image": ("IMAGE",), + "human_mask": ("MASK",) } } - RETURN_TYPES = ("IMAGE", "MASK") + RETURN_TYPES = ("IMAGE",) + FUNCTION = "edge_add" + CATEGORY = "facechain/crop" + + def edge_add(self, source_image, edge_add_image, human_mask): + np_source_image = tensor_to_np(source_image) + np_edge_add_origin_image = tensor_to_np(edge_add_image) + np_human_mask = mask_tensor_to_mask_np3(human_mask) + edge_add = np_source_image.astype(np.int16) - np_edge_add_origin_image.astype(np.int16) + edge_add = edge_add * (1 - np_human_mask) + result = Image.fromarray((np.clip(np_source_image.astype(np.int16) + edge_add.astype(np.int16), 0, 255)).astype(np.uint8)) + return (image_to_tensor(result),) + + +class FCReplaceByMask: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source_image": ("IMAGE",), + "replace_image": ("IMAGE",), + "mask": ("MASK",), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "replace_by_mask" + CATEGORY = "facechain/crop" + + def replace_by_mask(self, source_image, replace_image, mask): + np_source_image = tensor_to_np(source_image) + np_replace_image = tensor_to_np(replace_image) + np_mask = mask_tensor_to_mask_np3(mask) + result_np = np_source_image * np_mask + np_replace_image(1 - np_mask) + return (image_np_to_image_tensor(result_np),) + + +class FCCropAndPaste: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "template_image": ("IMAGE",), + "human_image": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "MASK") + RETURN_NAMES = ("crop_image", "mask", "invert_mask") FUNCTION = "crop_and_paste" CATEGORY = "facechain/crop" - def crop_and_paste(this, source_image, source_image_mask, source_box, source_five_point, target_image, target_five_point, use_warp=True): - source_image = tensor_to_img(source_image) - target_image = tensor_to_img(target_image) - source_image_mask = tensor_to_img(source_image_mask) - if use_warp: - source_five_point = np.reshape(source_five_point, [5, 2]) - np.array(source_box[:2]) - target_five_point = np.reshape(target_five_point, [5, 2]) - - Crop_Source_image = source_image.crop(np.int32(source_box)) - Crop_Source_image_mask = source_image_mask.crop(np.int32(source_box)) - source_five_point, target_five_point = np.array(source_five_point), np.array(target_five_point) - - tform = transform.SimilarityTransform() - tform.estimate(source_five_point, target_five_point) - M = tform.params[0:2, :] - - warped = cv2.warpAffine(np.array(Crop_Source_image), M, np.shape(target_image)[:2][::-1], borderValue=0.0) - warped_mask = cv2.warpAffine(np.array(Crop_Source_image_mask), M, np.shape(target_image)[:2][::-1], borderValue=0.0) - - mask = np.float32(warped_mask == 0) - output = mask * np.float32(target_image) + (1 - mask) * np.float32(warped) - else: - mask = np.float32(np.array(source_image_mask) == 0) - output = mask * np.float32(target_image) + (1 - mask) * np.float32(source_image) - return image_np_to_image_tensor(output), mask_np3_to_mask_tensor(mask) + def crop_and_paste(this, template_image, human_image): + pil_template_image = tensor_to_img(template_image) + pil_human_image = tensor_to_img(human_image) + _, _, _, template_five_point = facechain_detect_crop(pil_template_image, 0, 1.1, 'normal') + _, _, human_box, human_five_point = facechain_detect_crop(pil_human_image, 0, 1.5, 'normal') + _, human_mask, _ = segment(pil_human_image) + human_mask = np.expand_dims((human_mask * 255).astype(np.uint8), axis=2) + human_mask = np.concatenate([human_mask, human_mask, human_mask], axis=2) + pil_human_mask = Image.fromarray(human_mask) + output, mask = crop_and_paste(pil_human_image, pil_human_mask, pil_template_image, human_five_point, template_five_point, human_box) + return image_np_to_image_tensor(output), mask_np3_to_mask_tensor(1 - mask), mask_np3_to_mask_tensor(mask) class FCMaskOP: @@ -173,7 +233,7 @@ class FCMaskOP: return { "required": { "mask": ("MASK",), - "method": (["concatenate"],), + "method": (["expand_dims", "concatenate"],), } } @@ -186,4 +246,6 @@ class FCMaskOP: result = None if method == "concatenate": result = np.concatenate([mask, mask, mask], axis=2) + elif method == "expand_dims": + result = np.expand_dims(mask, axis=2) return (mask_np3_to_mask_tensor(result),) diff --git a/facechain/utils/batch_covert_utils.py b/facechain/utils/batch_covert_utils.py new file mode 100644 index 0000000..59b7da1 --- /dev/null +++ b/facechain/utils/batch_covert_utils.py @@ -0,0 +1,11 @@ +import numpy as np +from PIL import Image + + +def tensors_to_imgs(tensors): + images = [] + for image in tensors: + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)).convert("RGB") + images.append(image) + return images diff --git a/facechain/utils/convert_utils.py b/facechain/utils/convert_utils.py index 7eae6d9..fc48f97 100644 --- a/facechain/utils/convert_utils.py +++ b/facechain/utils/convert_utils.py @@ -1,8 +1,10 @@ +import cv2 import numpy as np import torch from PIL import ImageOps from PIL import Image + def image_to_tensor(input): i = ImageOps.exif_transpose(input) image = i.convert("RGB") @@ -10,18 +12,21 @@ def image_to_tensor(input): tensor = torch.from_numpy(image)[None,] return tensor + def image_to_np(input): i = ImageOps.exif_transpose(input) image = i.convert("RGB") image_np = np.array(image).astype(np.uint8) return image_np + def tensor_to_np(image): image = image[0] i = 255. * image.cpu().numpy() result = np.clip(i, 0, 255).astype(np.uint8) return result + def img_to_mask(input): i = ImageOps.exif_transpose(input) image = i.convert("RGB") @@ -29,23 +34,30 @@ def img_to_mask(input): mask_tensor = torch.from_numpy(new_np).permute(2, 0, 1)[0:1, :, :] return mask_tensor + def image_np_to_image_tensor(input): image = input.astype(np.float32) / 255.0 tensor = torch.from_numpy(image)[None,] return tensor + def mask_np2_to_mask_tensor(input): image = input.astype(np.float32) tensor = torch.from_numpy(image)[None,] return tensor + + def mask_np3_to_mask_tensor(input): image = input.astype(np.float32) tensor = torch.from_numpy(image).permute(2, 0, 1)[0:1, :, :] return tensor + + def mask_tensor_to_mask_np3(input): result = input.permute(1, 2, 0).cpu().numpy() return result + def tensor_to_img(image): image = image[0] i = 255. * image.cpu().numpy()