nodes/THtianhao-ComfyUI-FaceChain
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update some node

Browse Source
This commit is contained in:
tototianhao
2024-01-07 18:45:12 +08:00
parent a54881ac60
commit bd994dbcd5
6 changed files with 204 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+12 -5
View File
@@ -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")
facechain/common/model_processor.py
+55 -18
View File
@@ -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
facechain/model_holder.py
-1
View File
@@ -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
facechain/nodes.py
+114 -52
View File
@@ -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),)
facechain/utils/batch_covert_utils.py
+11
View File
@@ -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
facechain/utils/convert_utils.py
+12
View File
@@ -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()