This commit is contained in:
ultimatech
2023-11-24 08:25:32 +08:00
3 changed files with 122 additions and 13 deletions
+5
View File
@@ -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
View File
@@ -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),)
+5 -1
View File
@@ -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()