332 lines
11 KiB
Python
332 lines
11 KiB
Python
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
from modelscope.outputs import OutputKeys
|
|
from modelscope.pipelines import pipeline
|
|
from modelscope.utils.constant import Tasks
|
|
from .face_process_utils import call_face_crop, color_transfer, Face_Skin
|
|
from protrait.img_utils import img_to_tensor, tensor_to_img, tensor_to_np, np_to_tensor, np_to_mask, img_to_mask
|
|
from .config import models_path
|
|
|
|
# import pydevd_pycharm
|
|
# pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
|
|
|
|
class RetainFace:
|
|
def __init__(self):
|
|
self.retinaface_detection = pipeline(Tasks.face_detection, 'damo/cv_resnet50_face-detection_retinaface', model_revision='v2.0.2')
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image": ("IMAGE",),
|
|
"multi_user_facecrop_ratio": ("FLOAT", {"default": 1, "min": 0, "max": 10, "step": 0.1})
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "BOX")
|
|
RETURN_NAMES = ("crop_image", "crop_mask", "crop_box")
|
|
FUNCTION = "retain_face"
|
|
CATEGORY = "protrait/model"
|
|
|
|
def retain_face(self, image, multi_user_facecrop_ratio):
|
|
np_image = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
|
image = Image.fromarray(np_image)
|
|
retinaface_boxes, retinaface_keypoints, retinaface_masks, retinaface_tensor = call_face_crop(self.retinaface_detection, image, multi_user_facecrop_ratio)
|
|
crop_image = image.crop(retinaface_boxes[0])
|
|
return (img_to_tensor(crop_image), retinaface_tensor, retinaface_boxes[0])
|
|
|
|
class FaceFusionPM:
|
|
|
|
def __init__(self):
|
|
self.image_face_fusion = pipeline(Tasks.image_face_fusion, model='damo/cv_unet-image-face-fusion_damo', model_revision='v1.3')
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image": ("IMAGE",),
|
|
"user_image": ("IMAGE",),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "img_face_fusion"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def img_face_fusion(self, image, user_image):
|
|
image = tensor_to_img(image)
|
|
user_image = tensor_to_img(user_image)
|
|
fusion_image = self.image_face_fusion(dict(template=image, user=user_image))[
|
|
OutputKeys.OUTPUT_IMG]
|
|
# swap_face(target_img=output_image, source_img=roop_image, model="inswapper_128.onnx", upscale_options=UpscaleOptions())
|
|
fusion_image = Image.fromarray(cv2.cvtColor(fusion_image, cv2.COLOR_BGR2RGB))
|
|
return (img_to_tensor(fusion_image),)
|
|
|
|
class RatioMerge2Image:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
"fusion_rate": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.1})
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "image_ratio_merge"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def image_ratio_merge(self, image1, image2, fusion_rate):
|
|
rate_fusion_image = image1 * (1 - fusion_rate) + image2 * fusion_rate
|
|
return (rate_fusion_image,)
|
|
|
|
class ReplaceBoxImg:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"origin_image": ("IMAGE",),
|
|
"box_area": ("BOX",),
|
|
"replace_image": ("IMAGE",),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "replace_box_image"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def replace_box_image(self, origin_image, box_area, replace_image):
|
|
origin_image[:, box_area[1]:box_area[3], box_area[0]:box_area[2], :] = replace_image
|
|
return (origin_image,)
|
|
|
|
class MaskMerge2Image:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "image_mask_merge"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def image_mask_merge(self, image1, image2, mask, box=None):
|
|
mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
|
|
image1 = image1 * mask + image2 * (1 - mask)
|
|
return (image1,)
|
|
|
|
class ExpandMaskFaceWidth:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"mask": ("MASK",),
|
|
"box": ("BOX",),
|
|
"expand_width": ("FLOAT", {"default": 0.15, "min": 0, "max": 10, "step": 0.1})
|
|
}}
|
|
|
|
RETURN_TYPES = ("MASK", "BOX")
|
|
FUNCTION = "expand_mask_face_width"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def expand_mask_face_width(self, mask, box, expand_width):
|
|
h, w = mask.shape[1], mask.shape[2]
|
|
|
|
new_mask = mask.clone().zero_()
|
|
copy_box = np.copy(np.int32(box))
|
|
|
|
face_width = copy_box[2] - copy_box[0]
|
|
copy_box[0] = np.clip(np.array(copy_box[0], np.int32) - face_width * expand_width, 0, w - 1)
|
|
copy_box[2] = np.clip(np.array(copy_box[2], np.int32) + face_width * expand_width, 0, w - 1)
|
|
|
|
# get new input_mask
|
|
new_mask[0, copy_box[1]:copy_box[3], copy_box[0]:copy_box[2]] = 255
|
|
return (new_mask, copy_box)
|
|
|
|
class BoxCropImage:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"image": ("IMAGE",),
|
|
"box": ("BOX",), }
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("crop_image",)
|
|
FUNCTION = "box_crop_image"
|
|
CATEGORY = "protrait/model"
|
|
|
|
def box_crop_image(self, image, box):
|
|
image = image[:, box[1]:box[3], box[0]:box[2], :]
|
|
return (image,)
|
|
|
|
class ColorTransfer:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"transfer_from": ("IMAGE",),
|
|
"transfer_to": ("IMAGE",),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "color_transfer"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def color_transfer(self, transfer_from, transfer_to):
|
|
transfer_result = color_transfer(tensor_to_np(transfer_from), tensor_to_np(transfer_to)) # 进行颜色迁移
|
|
return (np_to_tensor(transfer_result),)
|
|
|
|
class FaceSkin:
|
|
def __init__(self):
|
|
self.retinaface_detection = pipeline(Tasks.face_detection, 'damo/cv_resnet50_face-detection_retinaface', model_revision='v2.0.2')
|
|
self.face_skin = Face_Skin(os.path.join(models_path, "face_skin.pth"))
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"image": ("IMAGE",), }
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "face_skin_mask"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def face_skin_mask(self, image):
|
|
face_skin_one = self.face_skin.detect(tensor_to_img(image), self.retinaface_detection, [1, 2, 3, 4, 5, 10, 12, 13])
|
|
return (face_skin_one,)
|
|
|
|
class MaskDilateErode:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"mask": ("MASK",), }
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "mask_dilate_erode"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def mask_dilate_erode(self, mask):
|
|
out_mask = Image.fromarray(np.uint8(cv2.dilate(tensor_to_np(mask), np.ones((96, 96), np.uint8), iterations=1) - cv2.erode(tensor_to_np(mask), np.ones((48, 48), np.uint8), iterations=1)))
|
|
return (img_to_mask(out_mask),)
|
|
|
|
class SkinRetouching:
|
|
|
|
def __init__(self):
|
|
self.skin_retouching = pipeline('skin-retouching-torch', model='damo/cv_unet_skin_retouching_torch', model_revision='v1.0.2')
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"image": ("IMAGE",)}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "skin_retouching_pass"
|
|
CATEGORY = "protrait/model"
|
|
|
|
def skin_retouching_pass(self, image):
|
|
output_image = cv2.cvtColor(self.skin_retouching(tensor_to_img(image))[OutputKeys.OUTPUT_IMG], cv2.COLOR_BGR2RGB)
|
|
return (np_to_tensor(output_image),)
|
|
|
|
class PortraitEnhancement:
|
|
|
|
def __init__(self):
|
|
self.portrait_enhancement = pipeline(Tasks.image_portrait_enhancement, model='damo/cv_gpen_image-portrait-enhancement', model_revision='v1.0.0')
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"image": ("IMAGE",), }
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "protrait_enhancement_pass"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def protrait_enhancement_pass(self, image):
|
|
output_image = cv2.cvtColor(self.portrait_enhancement(tensor_to_img(image))[OutputKeys.OUTPUT_IMG], cv2.COLOR_BGR2RGB)
|
|
return (np_to_tensor(output_image),)
|
|
|
|
class ImageScaleShort:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"image": ("IMAGE",),
|
|
"size": ("INT", {"default": 512, "min": 0, "max": 2048, "step": 1}),
|
|
"crop_face": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "image_scale_short"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def image_scale_short(self, image, size, crop_face):
|
|
input_image = tensor_to_img(image)
|
|
short_side = min(input_image.width, input_image.height)
|
|
resize = float(short_side / size)
|
|
new_size = (int(input_image.width // resize), int(input_image.height // resize))
|
|
input_image = input_image.resize(new_size, Image.Resampling.LANCZOS)
|
|
if crop_face:
|
|
new_width = int(np.shape(input_image)[1] // 32 * 32)
|
|
new_height = int(np.shape(input_image)[0] // 32 * 32)
|
|
input_image = input_image.resize([new_width, new_height], Image.Resampling.LANCZOS)
|
|
return (img_to_tensor(input_image),)
|
|
|
|
class ImageResizeTarget:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"image": ("IMAGE",),
|
|
"width": ("INT", {"default": 512, "min": 0, "max": 2048, "step": 1}),
|
|
"height": ("INT", {"default": 512, "min": 0, "max": 2048, "step": 1}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
FUNCTION = "image_resize_target"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def image_resize_target(self, image, width, height):
|
|
imagepi = tensor_to_img(image)
|
|
out = imagepi.resize([width, height], Image.Resampling.LANCZOS)
|
|
return (img_to_tensor(out),)
|
|
|
|
class GetImageInfo:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"image": ("IMAGE",),
|
|
}}
|
|
|
|
RETURN_TYPES = ("INT", "INT")
|
|
RETURN_NAMES = ("width", "height")
|
|
|
|
FUNCTION = "get_image_info"
|
|
|
|
CATEGORY = "protrait/model"
|
|
|
|
def get_image_info(self, image):
|
|
width = image.shape[2]
|
|
height = image.shape[1]
|
|
return (width, height)
|