nodes/ArtBot2023-CharacterFaceSwap
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

init

Browse Source
This commit is contained in:
ArtBot2023
2023-09-05 10:56:44 +08:00
commit 6a4b0eb614
8 changed files with 6799 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -0,0 +1,2 @@
__pycache__
.vscode
__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .nodes import NODE_CLASS_MAPPINGS
__all__ = ['NODE_CLASS_MAPPINGS']
blend.py
+133
View File
@@ -0,0 +1,133 @@
"""Implement Photoshop color blend
ref: https://stackoverflow.com/questions/12121393/please-explain-this-color-blending-mode-formula-so-i-can-replicate-it-in-php-ima
"""
from typing import Literal
import cv2
import numpy as np
R = 0.3
G = 0.59
B = 0.11
def rgb2hsy(image):
"""image is normalized to [0,1]
"""
image = np.minimum(np.maximum(image, 0), 1)
r, g, b = image[:,:,0], image[:,:,1], image[:,:,2]
h = np.zeros_like(r)
s = np.zeros_like(g)
y = R * r + G * g + B * b
mask_gray = (r == g) & (g == b)
s[mask_gray] = 0
h[mask_gray] = 0
mask_sector_0 = ((r >= g) & (g >= b)) & ~mask_gray # Sector 0: 0° - 60°
s[mask_sector_0] = r[mask_sector_0] - b[mask_sector_0]
h[mask_sector_0] = 60 * (g[mask_sector_0] - b[mask_sector_0]) / s[mask_sector_0]
mask_sector_1 = (g > r) & (r >= b)
s[mask_sector_1] = g[mask_sector_1] - b[mask_sector_1]
h[mask_sector_1] = 60 * (g[mask_sector_1] - r[mask_sector_1]) / s[mask_sector_1] + 60
mask_sector_2 = (g >= b) & (b > r)
s[mask_sector_2] = g[mask_sector_2] - r[mask_sector_2]
h[mask_sector_2] = 60 * (b[mask_sector_2] - r[mask_sector_2]) / s[mask_sector_2] + 120
mask_sector_3 = (b > g) & (g > r)
s[mask_sector_3] = b[mask_sector_3] - r[mask_sector_3]
h[mask_sector_3] = 60 * (b[mask_sector_3] - g[mask_sector_3]) / s[mask_sector_3] + 180
mask_sector_4 = (b > r) & (r >= g)
s[mask_sector_4] = b[mask_sector_4] - g[mask_sector_4]
h[mask_sector_4] = 60 * (r[mask_sector_4] - g[mask_sector_4]) / s[mask_sector_4] + 240
mask_sector_5 = ~(mask_gray | mask_sector_0 | mask_sector_1 | mask_sector_2 | mask_sector_3 | mask_sector_4)
s[mask_sector_5] = r[mask_sector_5] - g[mask_sector_5]
h[mask_sector_5] = 60 * (r[mask_sector_5] - b[mask_sector_5]) / s[mask_sector_5] + 300
hsy = np.zeros_like(image)
hsy[:,:,0] = h % 360
hsy[:,:,1] = np.minimum(np.maximum(s, 0), 1)
hsy[:,:,2] = np.minimum(np.maximum(y, 0), 1)
return hsy
def hsy2rgb(image):
h, s, y = image[:, :, 0], image[:, :, 1], image[:, :, 2]
h = h % 360
s = np.minimum(np.maximum(s, 0), 1)
y = np.minimum(np.maximum(y, 0), 1)
r = np.zeros_like(h)
g = np.zeros_like(h)
b = np.zeros_like(h)
k = np.zeros_like(h)
mask_sector_0 = (h >= 0) & (h < 60)
k[mask_sector_0] = s[mask_sector_0] * h[mask_sector_0] / 60
b[mask_sector_0] = y[mask_sector_0] - R * s[mask_sector_0] - G * k[mask_sector_0]
r[mask_sector_0] = b[mask_sector_0] + s[mask_sector_0]
g[mask_sector_0] = b[mask_sector_0] + k[mask_sector_0]
mask_sector_1 = (h >= 60) & (h < 120)
k[mask_sector_1] = s[mask_sector_1] * (h[mask_sector_1] - 60) / 60
g[mask_sector_1] = y[mask_sector_1] + B * s[mask_sector_1] + R * k[mask_sector_1]
b[mask_sector_1] = g[mask_sector_1] - s[mask_sector_1]
r[mask_sector_1] = g[mask_sector_1] - k[mask_sector_1]
mask_sector_2 = (h >= 120) & (h < 180)
k[mask_sector_2] = s[mask_sector_2] * (h[mask_sector_2] - 120) / 60
r[mask_sector_2] = y[mask_sector_2] - G * s[mask_sector_2] - B * k[mask_sector_2]
g[mask_sector_2] = r[mask_sector_2] + s[mask_sector_2]
b[mask_sector_2] = r[mask_sector_2] + k[mask_sector_2]
mask_sector_3 = (h >= 180) & (h < 240)
k[mask_sector_3] = s[mask_sector_3] * (h[mask_sector_3] - 180) / 60
b[mask_sector_3] = y[mask_sector_3] + R * s[mask_sector_3] + G * k[mask_sector_3]
r[mask_sector_3] = b[mask_sector_3] - s[mask_sector_3]
g[mask_sector_3] = b[mask_sector_3] - k[mask_sector_3]
mask_sector_4 = (h >= 240) & (h < 300)
k[mask_sector_4] = s[mask_sector_4] * (h[mask_sector_4] - 240) / 60
g[mask_sector_4] = y[mask_sector_4] - B * s[mask_sector_4] - R * k[mask_sector_4]
b[mask_sector_4] = g[mask_sector_4] + s[mask_sector_4]
r[mask_sector_4] = g[mask_sector_4] + k[mask_sector_4]
mask_sector_5 = h >= 300
k[mask_sector_5] = s[mask_sector_5] * (h[mask_sector_5] - 300) / 60
r[mask_sector_5] = y[mask_sector_5] + G * s[mask_sector_5] + B * k[mask_sector_5]
g[mask_sector_5] = r[mask_sector_5] - s[mask_sector_5]
b[mask_sector_5] = r[mask_sector_5] - k[mask_sector_5]
return np.minimum(np.maximum(np.stack([r, g, b], axis=-1), 0), 1)
def color_blend(base_image, blend_image, mode: Literal["Hue", "Saturation", "Color", "Luminosity"]):
"""
Args:
base_image (cv2)
blend_image (cv2)
Return:
cv2
"""
base_image = cv2.cvtColor(base_image, cv2.COLOR_BGR2RGB).astype(np.float32)/255
blend_image = cv2.cvtColor(blend_image, cv2.COLOR_BGR2RGB).astype(np.float32)/255
# Convert to HLS color space using OpenCV
hsy_base = rgb2hsy(base_image)
hsy_blend = rgb2hsy(blend_image)
if mode=="Hue":
hsy_out = np.stack([hsy_blend[:,:,0], hsy_base[:,:,1], hsy_base[:,:,2]], axis=-1)
elif mode=="Saturation":
hsy_out = np.stack([hsy_base[:,:,0], hsy_blend[:,:,1], hsy_base[:,:,2]], axis=-1)
elif mode=="Color":
hsy_out = np.stack([hsy_blend[:,:,0], hsy_blend[:,:,1], hsy_base[:,:,2]], axis=-1)
elif mode=="Luminosity":
hsy_out = np.stack([hsy_base[:,:,0], hsy_base[:,:,1], hsy_blend[:,:,2]], axis=-1)
else: assert False, f"{mode} is not a valid mode"
rgb_out = hsy2rgb(hsy_out)
return cv2.cvtColor((rgb_out*255).astype(np.uint8), cv2.COLOR_RGB2BGR)
nodes.py
+448
View File
@@ -0,0 +1,448 @@
from typing import Tuple, List, Literal
import os
import comfy.samplers
from facexlib.detection import RetinaFace
from facexlib.parsing import BiSeNet
import torch
import numpy as np
import cv2
from .utils import models_dir, tensor2pil, pil2tensor, tensor2cv, cv2tensor, hex2bgr, BBox
# copy from SeargeSDXL
class GenerationParameterInput:
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"image_width": ("INT", {"default": 512, "min": 0, "max": 1024, "step": 8}),
"image_height": ("INT", {"default": 512, "min": 0, "max": 1024, "step": 8}),
"steps": ("INT", {"default": 20, "min": 0, "max": 200}),
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 30.0, "step": 0.5}),
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "dpmpp_2m"}),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "karras"}),
"denoise": ("FLOAT", {"default": 1, "min": 0, "max": 1})
},
"optional": {
"parameters": ("PARAMETERS", ),
},
}
RETURN_TYPES = ("PARAMETERS", )
RETURN_NAMES = ("parameters", )
FUNCTION = "mux"
CATEGORY = "ArtBot2023"
def mux(self, seed, image_width, image_height, steps, cfg, sampler_name, scheduler, denoise, parameters={}):
parameters["seed"] = seed
parameters["image_width"] = image_width
parameters["image_height"] = image_height
parameters["steps"] = steps
parameters["cfg"] = cfg
parameters["sampler_name"] = sampler_name
parameters["scheduler"] = scheduler
parameters["denoise"] = denoise
return (parameters, )
class GenertaionParameterOutput:
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"parameters": ("PARAMETERS", ),
},
}
RETURN_TYPES = ("PARAMETERS", "INT", "INT", "INT", "INT", "FLOAT", comfy.samplers.KSampler.SAMPLERS, comfy.samplers.KSampler.SCHEDULERS, "FLOAT", )
RETURN_NAMES = ("parameters", "seed", "image_width", "image_height", "steps", "cfg", "sampler_name", "scheduler", "denoise", )
FUNCTION = "demux"
CATEGORY = "ArtBot2023"
def demux(self, parameters):
seed = parameters["seed"]
image_width = parameters["image_width"]
image_height = parameters["image_height"]
steps = parameters["steps"]
cfg = parameters["cfg"]
sampler_name = parameters["sampler_name"]
scheduler = parameters["scheduler"]
denoise = parameters["denoise"]
return (parameters, seed, image_width, image_height, steps, cfg, sampler_name, scheduler, denoise, )
class LoadRetinaFace:
models_dir = os.path.join(models_dir, 'facexlib')
@classmethod
def INPUT_TYPES(cls):
return {"required":{}}
RETURN_TYPES = ("RETINAFACE", )
RETURN_NAMES = ("MODEL", )
FUNCTION = "load"
CATEGORY = "ArtBot2023"
def load(self):
from facexlib.detection import init_detection_model
return (init_detection_model("retinaface_resnet50", model_rootpath=self.models_dir), )
class CropFace:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("RETINAFACE", ),
"image": ("IMAGE", ),
"confidence": ("FLOAT", {"default": 0.8, "min": 0, "max": 1}),
"margin": ("INT", {"default": 32}),
}
}
RETURN_TYPES = (
"IMAGE", "IMAGE", "BBOX"
)
RETURN_NAMES = ("face_image", "preview", "bbox")
FUNCTION = "crop"
CATEGORY = "ArtBot2023"
def crop(self, model: RetinaFace, image: torch.Tensor, confidence: float, margin: int):
with torch.no_grad():
# model receives bgr uint8 format
# bboxes: list of [x0, y0, x1, y1, confidence_score, five points (x, y)]
bboxes = model.detect_faces(tensor2cv(image), confidence)
if (len(bboxes)==0):
print("no face detected")
return np.zeros((4,)), 0, image
detection_preview = self.visualize_detection(tensor2cv(image), bboxes)
bboxes = [self.add_margin_and_make_square((int(min(x0, x1)), int(min(y0,y1)), int(abs(x1-x0)), int(abs(y1-y0))), margin, img_width=image.shape[2], img_height=image.shape[1]) for (x0, y0, x1, y1, *_) in bboxes]
detection_preview = self.visualize_margin(detection_preview, bboxes)
faces = self.crop_faces(bboxes, image)
# scaled_faces = self.scale_faces(faces, size)
return faces[0].unsqueeze(0), cv2tensor(detection_preview), bboxes[0]
def crop_faces(self, bboxes: List[BBox], image: torch.Tensor):
"""
Returns: list of Tensor[h,w,c]
"""
return [image[0, y:y+h, x:x+w, :] for (x,y,w,h) in bboxes]
def scale_faces(self, faces: List[torch.Tensor], size: int, upscaler: Literal["linear"]="linear"):
"""
Args:
faces: list of Tensor[h,w,c]
"""
scaled_faces: List[torch.Tensor] = []
for face in faces:
# Change the layout to [batch, channel, height, width]
face = face.permute(2, 0, 1).unsqueeze(0)
# Perform the interpolation
if upscaler == "linear":
scaled_face = torch.nn.functional.interpolate(face, size=(size, size), mode="bilinear", align_corners=True)
elif upscaler == "nearest":
scaled_face = torch.nn.functional.interpolate(face, size=(size, size), mode="nearest")
else:
raise ValueError(f"Invalid upscaler: {upscaler}")
# Change the layout back to [height, width, channel] and remove batch dimension
scaled_face = scaled_face.squeeze(0).permute(1, 2, 0)
scaled_faces.append(scaled_face)
return scaled_faces
def visualize_margin(self, img, bboxes):
img = np.copy(img)
for bbox in bboxes:
x,y,w,h = bbox
cv2.rectangle(img, (x,y), (x+w, y+h), hex2bgr("#710193"), 2)
return img
def visualize_detection(self, img, bboxes_and_landmarks):
"""
Args:
img (np.ndarray): bgr
Returns:
img: bgr
"""
img = np.copy(img)
for b in bboxes_and_landmarks:
# confidence
cv2.putText(img, f'{b[4]:.4f}', (int(b[0]), int(b[1] + 12)), cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))
# bounding boxes
b = list(map(int, b))
cv2.rectangle(img, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2)
# landmarks (for retinaface)
cv2.circle(img, (b[5], b[6]), 1, (0, 0, 255), 4)
cv2.circle(img, (b[7], b[8]), 1, (0, 255, 255), 4)
cv2.circle(img, (b[9], b[10]), 1, (255, 0, 255), 4)
cv2.circle(img, (b[11], b[12]), 1, (0, 255, 0), 4)
cv2.circle(img, (b[13], b[14]), 1, (255, 0, 0), 4)
return img
def add_margin_and_make_square(self, bbox: BBox, margin: int, img_width: int, img_height: int):
x, y, w, h = map(lambda x: int(x), bbox) # x, y are the coordinates of the top-left corner of the bounding box
# Calculate margin
margin_w = margin
margin_h = margin
# Add margin to the bounding box, ensuring it doesn't go out of the image boundaries
x = max(0, x - margin_w)
y = max(0, y - margin_h)
w = min(img_width - x, w + 2 * margin_w)
h= min(img_height - y, h + 2 * margin_h)
# Make the bounding box square while keeping the center the same
cx, cy = x + w // 2, y + h // 2 # Calculate the center of the original bounding box
max_side = max(w, h)
x = max(0, cx - max_side // 2)
y = max(0, cy - max_side // 2)
w = h = min(max_side, img_width - x, img_height - y) # Ensure the bounding box is within the image boundaries
return int(x), int(y), int(w), int(h)
class UncropFace:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", ),
"bbox": ("BBOX", ),
"face": ("IMAGE", ),
"mask": ("MASK", )
}
}
RETURN_TYPES = ("IMAGE", )
CATEGORY = "ArtBot2023"
FUNCTION = "uncrop"
def uncrop(self, image: torch.Tensor, bbox: BBox, face: torch.Tensor, mask: torch.Tensor):
bbox_face, bbox_mask = self.scale_face(face.squeeze(), mask, bbox[2])
cv2.imwrite("bbox_face.png", tensor2cv(bbox_face))
cv2.imwrite("bbox_mask.png", tensor2cv(bbox_mask))
image_apply_face = self.weighted_sum(image.squeeze(), bbox, bbox_face, bbox_mask)
return (image_apply_face.unsqueeze(0), )
def scale_face(self, face: torch.Tensor, mask: torch.Tensor, size):
"""
Args:
face (torch.Tensor): [h,w,c]
mask (torch.Tensor): [h,w]
"""
scaled_faces: List[torch.Tensor] = []
for face in [face, mask.unsqueeze(-1)]:
# Change the layout to [batch, channel, height, width]
face = face.permute(2, 0, 1).unsqueeze(0)
scaled_face = torch.nn.functional.interpolate(face, size=(size, size), mode="bilinear", align_corners=True)
# Change the layout back to [height, width, channel] and remove batch dimension
scaled_face = scaled_face.squeeze(0).permute(1, 2, 0)
scaled_faces.append(scaled_face)
return scaled_faces[0], scaled_faces[1].squeeze()
def weighted_sum(self, image: torch.Tensor, bbox: BBox, face: torch.Tensor, mask: torch.Tensor)->torch.Tensor:
"""
Args:
image (torch.Tensor): [h_full,w_full,c]
bbox (BBox): [x,y,w,h]
face (torch.Tensor): [h,w,c]
mask (torch.Tensor): [h,w]
Returns:
torch.Tensor: same shape as image
"""
image = image.clone()
x,y,w,h = bbox
mask = mask.unsqueeze(-1)
image[y:y+h, x:x+w, :] = mask * face + image[y:y+h, x:x+w, :] * (1-mask)
return image
class LoadBisenet:
models_dir = os.path.join(models_dir, 'facexlib')
@classmethod
def INPUT_TYPES(cls):
return {"required":{}}
RETURN_TYPES = ("BISENET", )
FUNCTION = "load"
CATEGORY = "ArtBot2023"
def load(self):
from facexlib.parsing import init_parsing_model
return (init_parsing_model("bisenet", model_rootpath=self.models_dir), )
class SegFace:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("BISENET", ),
"image": ("IMAGE", ),
"expand": ("INT", {"min": 0}),
"include_hair": (["enable", "disable"], {"default": "disable"}),
"include_neck": (["enable", "disable"], {"default": "disable"}),
}
}
RETURN_TYPES = (
"IMAGE", "MASK"
)
RETURN_NAMES = ("image", "mask")
FUNCTION = "segment"
CATEGORY = "ArtBot2023"
# labels: 0 'background'
# 1 'skin', 2 'l_brow', 3 'r_brow', 4 'l_eye', 5 'r_eye',
# 6 'eye_g', 7 'l_ear', 8 'r_ear', 9 'ear_r', 10 'nose',
# 11 'mouth', 12 'u_lip', 13 'l_lip', 14 'neck', 15 'neck_l',
# 16 'cloth', 17 'hair', 18 'hat'
annotation_name = ['background',
'skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye',
'eye_g', 'l_ear', 'r_ear', 'ear_r', 'nose',
'mouth', 'u_lip', 'l_lip', 'neck', 'neck_l',
'cloth', 'hair', 'hat']
annotation_color = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 0, 85], [255, 0, 170], [0, 255, 0], [85, 255, 0],
[170, 255, 0], [0, 255, 85], [0, 255, 170], [0, 0, 255], [85, 0, 255], [170, 0, 255], [0, 85, 255],
[0, 170, 255], [255, 255, 0], [255, 255, 85], [255, 255, 170], [255, 0, 255], [255, 85, 255],
[255, 170, 255], [0, 255, 255], [85, 255, 255], [170, 255, 255]]
def segment(self, model: BiSeNet, image: torch.Tensor, expand, include_hair, include_neck):
image = image.squeeze().permute(2,0,1).unsqueeze(0).flip([1]) # shape [1, c, h, w], rgb2bgr
with torch.no_grad():
from torchvision.transforms.functional import normalize
out = model(normalize(image, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]).cuda())[0] # shape [1, 19, h, w]
annotation = out.squeeze().cpu().numpy().argmax(0)
mask = self.get_mask(annotation, expand, include_hair, include_neck)
masked_face = mask.unsqueeze(0).unsqueeze(-1) * image.squeeze().permute(1,2,0).unsqueeze(0).flip([3]) # shape [1, h, w, c], bgr2rgb
return masked_face, mask
def get_mask(self, annotation, expand, include_hair, include_neck):
face_inds = list(range(1,14))
hair_ind = self.annotation_name.index("hair")
neck_ind = self.annotation_name.index("neck")
target_inds = face_inds
if include_hair == "enable": target_inds.append(hair_ind)
if include_neck == "enable": target_inds.append(neck_ind)
mask = np.zeros_like(annotation, dtype=np.float32)
for ind in target_inds:
mask[annotation==ind] = 1
blurred_mask = cv2.dilate(mask, np.ones((expand, expand), np.float32), iterations=1)
return torch.from_numpy(blurred_mask)
class ImageFullBBox:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", )
}
}
RETURN_TYPES = ("BBOX", )
FUNCTION = "bbox"
CATEGORY = "ArtBot2023"
def bbox(self, image: torch.Tensor):
image = image.squeeze()
return ((0,0,image.shape[1],image.shape[0]), )
class ColorBlend:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"blend_image": ("IMAGE", ),
"base_image": ("IMAGE", ),
"mode": (["Hue", "Saturation", "Color", "Luminosity"], )
}
}
RETURN_TYPES = ("IMAGE", )
FUNCTION = "blend"
CATEGORY = "ArtBot2023"
def blend(self, blend_image: torch.Tensor, base_image: torch.Tensor, mode: Literal["Hue", "Saturation", "Color", "Luminosity"]):
from .blend import color_blend
return (cv2tensor(color_blend(base_image=tensor2cv(base_image), blend_image=tensor2cv(blend_image), mode=mode)), )
class ExcludeFacialFeature:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"face": ("IMAGE", ),
"model": ("BISENET", ),
"image": ("IMAGE", ),
"expand": ("INT", {"min": 0})
}
}
RETURN_TYPES = ("IMAGE", )
FUNCTION = "exclude"
CATEGORY = "ArtBot2023"
annotation_name = ['background',
'skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye',
'eye_g', 'l_ear', 'r_ear', 'ear_r', 'nose',
'mouth', 'u_lip', 'l_lip', 'neck', 'neck_l',
'cloth', 'hair', 'hat']
def exclude(self, face: torch.Tensor, model: BiSeNet, image: torch.Tensor, expand: int):
face = face.squeeze().permute(2,0,1).unsqueeze(0).flip([1]) # shape [1, c, h, w], rgb2bgr
with torch.no_grad():
from torchvision.transforms.functional import normalize
out = model(normalize(face, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]).cuda())[0] # shape [1, 19, h, w]
annotation = out.squeeze().cpu().numpy().argmax(0)
# mask = self.get_mask_facial_feature(annotation, expand)
# image = image * (1-mask).unsqueeze(0).unsqueeze(-1)
mask = self.get_mask_hair(annotation, expand)
image = image * mask.unsqueeze(0).unsqueeze(-1)
return (image, )
def get_mask_hair(self, annotation, expand):
hair_ind = self.annotation_name.index('hair')
mask = np.zeros_like(annotation, dtype=np.float32)
mask[annotation==hair_ind] = 1
blurred_mask = cv2.dilate(mask, np.ones((expand, expand), np.float32), iterations=1)
return torch.from_numpy(blurred_mask)
def get_mask_facial_feature(self, annotation, expand):
facial_feature_inds = list(range(2,14))
mask = np.zeros_like(annotation, dtype=np.float32)
for ind in facial_feature_inds:
mask[annotation==ind] = 1
blurred_mask = cv2.dilate(mask, np.ones((expand, expand), np.float32), iterations=1)
return torch.from_numpy(blurred_mask)
class MaskContour:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mask": ("MASK", ),
}
}
RETURN_TYPES = ("MASK", )
FUNCTION = "find_contour"
CATEGORY = "ArtBot2023"
def find_contour(self, mask: torch.Tensor):
mask_np: np.ndarray = mask.squeeze().cpu().numpy().astype('uint8')
contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
mask_contour = np.zeros((mask_np.shape[0], mask_np.shape[1]), dtype=np.float32)
cv2.drawContours(mask_contour, contours, -1, (1, ), 1)
return (torch.from_numpy(mask_contour), )
NODE_CLASS_MAPPINGS = {
"Generation Parameter Input": GenerationParameterInput,
"Generation Parameter Output": GenertaionParameterOutput,
"Load RetinaFace": LoadRetinaFace,
"Load BiseNet": LoadBisenet,
"Uncrop Face": UncropFace,
"Crop Face": CropFace,
"Segment Face": SegFace,
"Image Full BBox": ImageFullBBox,
"Color Blend": ColorBlend,
"Exclude Facial Feature": ExcludeFacialFeature,
"Mask Contour": MaskContour
}
tifa_text_gen.json
+3953
View File
File diff suppressed because it is too large Load Diff
tmp.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.6 KiB

utils.py
+55
View File
@@ -0,0 +1,55 @@
import torch
import numpy as np
import folder_paths as comfy_paths
import comfy
from PIL import Image
import hashlib
import cv2
from typing import Tuple
BBox = Tuple[int, int, int, int]
models_dir = comfy_paths.models_dir
# Tensor to PIL
def tensor2pil(image):
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
# PIL to Tensor
def pil2tensor(image):
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
# PIL Hex
def pil2hex(image):
return hashlib.sha256(np.array(tensor2pil(image)).astype(np.uint16).tobytes()).hexdigest()
# PIL to Mask
def pil2mask(image):
image_np = np.array(image.convert("L")).astype(np.float32) / 255.0
mask = torch.from_numpy(image_np)
return 1.0 - mask
# Mask to PIL
def mask2pil(mask):
if mask.ndim > 2:
mask = mask.squeeze(0)
mask_np = mask.cpu().numpy().astype('uint8')
mask_pil = Image.fromarray(mask_np, mode="L")
return mask_pil
# Tensor to cv2
def tensor2cv(image):
image_np = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
return cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
# cv2 to Tensor
def cv2tensor(image):
image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
return torch.from_numpy(image_np).unsqueeze(0)
def hex2rgb(hex_color: str):
hex_color = hex_color.lstrip('#')
return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
def hex2bgr(hex_color):
return hex2rgb(hex_color)[::-1]
workflows/CharacterFaceSwap.json
+2205
View File
File diff suppressed because it is too large Load Diff