Merge branch 'main' of https://github.com/THtianhao/ComfyUI-FaceChain
This commit is contained in:
@@ -17,6 +17,9 @@ NODE_CLASS_MAPPINGS = {
|
||||
"FC_Segment": FCSegment,
|
||||
"FC_ReplaceImage": FCReplaceImage,
|
||||
"FC_CropBottom": FCCropBottom,
|
||||
"FC_CropFace": FCCropFace,
|
||||
"FC_CropAndPaste": FCCropAndPaste,
|
||||
"FC_MaskOP": FCMaskOP,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FC_FaceFusion": "FC FaceFusion",
|
||||
@@ -25,6 +28,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FC_CropMask": "FC CropMask",
|
||||
"FC_ReplaceImage": "FC ReplaceImage",
|
||||
"FC_CropBottom": "FC CropBottom",
|
||||
"FC_CropAndPaste": "FC CropAndPaste",
|
||||
"FC_MaskOP": "FC MaskOP",
|
||||
|
||||
}
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
+112
-12
@@ -3,20 +3,11 @@ import json
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from diffusers import StableDiffusionPipeline, StableDiffusionControlNetPipeline, ControlNetModel, \
|
||||
UniPCMultistepScheduler
|
||||
from skimage import transform
|
||||
from modelscope.outputs import OutputKeys
|
||||
from modelscope.pipelines import pipeline
|
||||
from modelscope.utils.constant import Tasks
|
||||
from torch import multiprocessing
|
||||
from transformers import pipeline as tpipeline
|
||||
import pydevd_pycharm
|
||||
|
||||
# import pydevd_pycharm
|
||||
#
|
||||
# pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
|
||||
pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
|
||||
|
||||
from .model_holder import *
|
||||
from .utils.img_utils import *
|
||||
@@ -287,3 +278,112 @@ class FCCropBottom:
|
||||
source_image = tensor_to_img(source_image)
|
||||
crop_result = crop_bottom(source_image, width)
|
||||
return (img_to_tensor(crop_result),)
|
||||
|
||||
class FCCropFace:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"source_image": ("IMAGE",),
|
||||
"crop_ratio": ("FLOAT", {"default": 1.0, "min": 0, "max": 10, "step": 0.1})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "BOX", "KEY_POINT")
|
||||
FUNCTION = "face_crop"
|
||||
CATEGORY = "facechain/crop"
|
||||
|
||||
def face_crop(self, source_image, crop_ratio):
|
||||
source_image_pil = tensor_to_img(source_image)
|
||||
det_result = get_face_detection()(source_image_pil)
|
||||
bboxes = det_result['boxes']
|
||||
keypoints = det_result['keypoints']
|
||||
area = 0
|
||||
idx = 0
|
||||
for i in range(len(bboxes)):
|
||||
bbox = bboxes[i]
|
||||
area_tmp = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
|
||||
if area_tmp > area:
|
||||
area = area_tmp
|
||||
idx = i
|
||||
bbox = bboxes[idx]
|
||||
keypoint = keypoints[idx]
|
||||
points_array = np.zeros((5, 2))
|
||||
for k in range(5):
|
||||
points_array[k, 0] = keypoint[2 * k]
|
||||
points_array[k, 1] = keypoint[2 * k + 1]
|
||||
w, h = source_image_pil.size
|
||||
face_w = bbox[2] - bbox[0]
|
||||
face_h = bbox[3] - bbox[1]
|
||||
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)
|
||||
bbox[3] = np.clip(np.array(bbox[3], np.int32) + face_h * (crop_ratio - 1) / 2, 0, h - 1)
|
||||
bbox = np.array(bbox, np.int32)
|
||||
result_image = source_image[:, bbox[1]:bbox[3], bbox[0]:bbox[2], :]
|
||||
return result_image, bbox, points_array
|
||||
|
||||
class FCCropAndPaste:
|
||||
@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",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "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)
|
||||
|
||||
class FCMaskOP:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",),
|
||||
"method": (["concatenate"],),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "mask_op"
|
||||
CATEGORY = "facechain/mask"
|
||||
|
||||
def mask_op(self, mask, method):
|
||||
mask = mask_tensor_to_mask_np3(mask)
|
||||
result = None
|
||||
if method == "concatenate":
|
||||
result = np.concatenate([mask, mask, mask], axis=2)
|
||||
return (mask_np3_to_mask_tensor(result),)
|
||||
|
||||
@@ -23,7 +23,7 @@ def img_to_mask(input):
|
||||
mask_tensor = torch.from_numpy(new_np).permute(2, 0, 1)[0:1, :, :]
|
||||
return mask_tensor
|
||||
|
||||
def image_np2_to_mask_tensor(input):
|
||||
def image_np_to_image_tensor(input):
|
||||
image = input.astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(image)[None,]
|
||||
return tensor
|
||||
@@ -36,6 +36,10 @@ 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()
|
||||
|
||||
Reference in New Issue
Block a user