chflame163
3bf19861c6
Commit MaskBoxDetect node, which can automatically detect the position through the mask and output it to the composite node. Commit X,Y to Percent node to convert absolute coordinates to percent coordinates. Commit GaussianBlur node. Commit GetImageSize node.
346 lines
14 KiB
Python
346 lines
14 KiB
Python
'''Image process functions for ComfyUI nodes
|
|
by chflame https://github.com/chflame163
|
|
'''
|
|
import numpy as np
|
|
import torch
|
|
import scipy.ndimage
|
|
import cv2
|
|
from typing import Union, List
|
|
from PIL import Image, ImageFilter, ImageChops, ImageDraw
|
|
|
|
def log(message):
|
|
name = 'LayerStyle'
|
|
print(f"# 😺dzNodes: {name} -> {message}")
|
|
|
|
'''Converter'''
|
|
|
|
def cv22pil(cv2_img:np.ndarray) -> Image:
|
|
cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
|
|
return Image.fromarray(cv2_img)
|
|
|
|
def pil2cv2(pil_img:Image) -> np.array:
|
|
np_img_array = np.asarray(pil_img)
|
|
return cv2.cvtColor(np_img_array, cv2.COLOR_RGB2BGR)
|
|
|
|
def pil2tensor(image:Image) -> torch.Tensor:
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def np2tensor(img_np: Union[np.ndarray, List[np.ndarray]]) -> torch.Tensor:
|
|
if isinstance(img_np, list):
|
|
return torch.cat([np2tensor(img) for img in img_np], dim=0)
|
|
return torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def tensor2np(tensor: torch.Tensor) -> List[np.ndarray]:
|
|
if len(tensor.shape) == 3: # Single image
|
|
return np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8)
|
|
else: # Batch of images
|
|
return [np.clip(255.0 * t.cpu().numpy(), 0, 255).astype(np.uint8) for t in tensor]
|
|
|
|
def tensor2pil(t_image: torch.Tensor) -> Image:
|
|
return Image.fromarray(np.clip(255.0 * t_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
def image2mask(image:Image) -> torch.Tensor:
|
|
_image = image.convert('RGBA')
|
|
alpha = _image.split() [0]
|
|
bg = Image.new("L", _image.size)
|
|
_image = Image.merge('RGBA', (bg, bg, bg, alpha))
|
|
ret_mask = torch.tensor([pil2tensor(_image)[0, :, :, 3].tolist()])
|
|
return ret_mask
|
|
|
|
def mask2image(mask:torch.Tensor) -> Image:
|
|
masks = tensor2np(mask)
|
|
for m in masks:
|
|
_mask = Image.fromarray(m).convert("L")
|
|
_image = Image.new("RGBA", _mask.size, color='white')
|
|
_image = Image.composite(
|
|
_image, Image.new("RGBA", _mask.size, color='black'), _mask)
|
|
return _image
|
|
|
|
'''Image Functions'''
|
|
|
|
def shift_image(image:Image, distance_x:int, distance_y:int) -> Image:
|
|
bkcolor = (0, 0, 0)
|
|
width = image.width
|
|
height = image.height
|
|
ret_image = Image.new('RGB', size=(width, height), color=bkcolor)
|
|
for x in range(width):
|
|
for y in range(height):
|
|
if x > -distance_x and y > -distance_y: # 防止回转
|
|
if x + distance_x < width and y + distance_y < height: # 防止越界
|
|
pixel = image.getpixel((x + distance_x, y + distance_y))
|
|
ret_image.putpixel((x, y), pixel)
|
|
return ret_image
|
|
|
|
def chop_image(background_image:Image, layer_image:Image, blend_mode:str, opacity:int) -> Image:
|
|
ret_image = background_image
|
|
if blend_mode == 'normal':
|
|
ret_image = layer_image
|
|
if blend_mode == 'multply':
|
|
ret_image = ImageChops.multiply(background_image,layer_image)
|
|
if blend_mode == 'screen':
|
|
ret_image = ImageChops.screen(background_image, layer_image)
|
|
if blend_mode == 'add':
|
|
ret_image = ImageChops.add(background_image, layer_image, 1, 0)
|
|
if blend_mode == 'subtract':
|
|
ret_image = ImageChops.subtract(background_image, layer_image, 1, 0)
|
|
if blend_mode == 'difference':
|
|
ret_image = ImageChops.difference(background_image, layer_image)
|
|
if blend_mode == 'darker':
|
|
ret_image = ImageChops.darker(background_image, layer_image)
|
|
if blend_mode == 'lighter':
|
|
ret_image = ImageChops.lighter(background_image, layer_image)
|
|
# opacity
|
|
if opacity == 0:
|
|
ret_image = background_image
|
|
elif opacity < 100:
|
|
alpha = 1.0 - float(opacity) / 100
|
|
ret_image = Image.blend(ret_image, background_image, alpha)
|
|
return ret_image
|
|
|
|
def remove_background(image:Image, mask:Image, color:str) -> Image:
|
|
width = image.width
|
|
height = image.height
|
|
ret_image = Image.new('RGB', size=(width, height), color=color)
|
|
ret_image.paste(image, mask=mask)
|
|
return ret_image
|
|
|
|
def gaussian_blur(image:Image, radius:int) -> Image:
|
|
image = image.convert("RGBA")
|
|
ret_image = image.filter(ImageFilter.GaussianBlur(radius=radius))
|
|
return ret_image
|
|
|
|
def motion_blur(image:Image, angle:int, blur:int) -> Image:
|
|
angle += 45
|
|
blur *= 5
|
|
image = np.array(pil2cv2(image))
|
|
M = cv2.getRotationMatrix2D((blur / 2, blur / 2), angle, 1)
|
|
motion_blur_kernel = np.diag(np.ones(blur))
|
|
motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M, (blur, blur))
|
|
motion_blur_kernel = motion_blur_kernel / blur
|
|
blurred = cv2.filter2D(image, -1, motion_blur_kernel)
|
|
# convert to uint8
|
|
cv2.normalize(blurred, blurred, 0, 255, cv2.NORM_MINMAX)
|
|
blurred = np.array(blurred, dtype=np.uint8)
|
|
ret_image = cv22pil(blurred)
|
|
return ret_image
|
|
|
|
def __rotate_expand(image:Image, angle:float, SSAA:int=0, method:str="lanczos") -> Image:
|
|
images = pil2tensor(image)
|
|
expand = "true"
|
|
height, width = images[0, :, :, 0].shape
|
|
|
|
def rotate_tensor(tensor):
|
|
resize_sampler = Image.LANCZOS
|
|
rotate_sampler = Image.BICUBIC
|
|
if method == "bicubic":
|
|
resize_sampler = Image.BICUBIC
|
|
rotate_sampler = Image.BICUBIC
|
|
elif method == "hamming":
|
|
resize_sampler = Image.HAMMING
|
|
rotate_sampler = Image.BILINEAR
|
|
elif method == "bilinear":
|
|
resize_sampler = Image.BILINEAR
|
|
rotate_sampler = Image.BILINEAR
|
|
elif method == "box":
|
|
resize_sampler = Image.BOX
|
|
rotate_sampler = Image.NEAREST
|
|
elif method == "nearest":
|
|
resize_sampler = Image.NEAREST
|
|
rotate_sampler = Image.NEAREST
|
|
if SSAA > 1:
|
|
img = tensor.tensor_to_image()
|
|
img_us_scaled = img.resize((width * SSAA, height * SSAA), resize_sampler)
|
|
img_rotated = img_us_scaled.rotate(angle, rotate_sampler, expand == "true", fillcolor=(0, 0, 0, 0))
|
|
img_down_scaled = img_rotated.resize((img_rotated.width // SSAA, img_rotated.height // SSAA), resize_sampler)
|
|
result = img_down_scaled.image_to_tensor()
|
|
else:
|
|
img = tensor.tensor_to_image()
|
|
img_rotated = img.rotate(angle, rotate_sampler, expand == "true", fillcolor=(0, 0, 0, 0))
|
|
result = img_rotated.image_to_tensor()
|
|
return result
|
|
|
|
if angle == 0.0 or angle == 360.0:
|
|
return tensor2pil(images)
|
|
else:
|
|
rotated_tensor = torch.stack([rotate_tensor(images[i]) for i in range(len(images))])
|
|
return tensor2pil(rotated_tensor).convert('RGB')
|
|
|
|
def image_rotate_extend_with_alpha(image:Image, angle:float, alpha:Image=None, method:str="lanczos", SSAA:int=0) -> tuple:
|
|
_image = __rotate_expand(image.convert('RGB'), angle, SSAA, method)
|
|
if angle is not None:
|
|
_alpha = __rotate_expand(alpha.convert('RGB'), angle, SSAA, method)
|
|
ret_image = RGB2RGBA(_image, _alpha)
|
|
else:
|
|
ret_image = _image
|
|
return (_image, _alpha, ret_image)
|
|
|
|
|
|
def create_gradient(start_color_inhex:str, end_color_inhex:str, width:int, height:int) -> Image:
|
|
start_color = Hex_to_RGB(start_color_inhex)
|
|
end_color = Hex_to_RGB(end_color_inhex)
|
|
ret_image = Image.new("RGB", (width, height), start_color)
|
|
draw = ImageDraw.Draw(ret_image)
|
|
for i in range(height):
|
|
R = int(start_color[0] * (height - i) / height + end_color[0] * i / height)
|
|
G = int(start_color[1] * (height - i) / height + end_color[1] * i / height)
|
|
B = int(start_color[2] * (height - i) / height + end_color[2] * i / height)
|
|
color = (R, G, B)
|
|
draw.line((0, i, width, i), fill=color)
|
|
return ret_image
|
|
|
|
def gradint(start_color_inhex:str, end_color_inhex:str, width:int, height:int, angle:float, ) -> Image:
|
|
radius = int((width + height) / 4)
|
|
g = create_gradient(start_color_inhex, end_color_inhex, radius, radius)
|
|
_canvas = Image.new('RGB', size=(radius, radius*3), color=start_color_inhex)
|
|
top = Image.new('RGB', size=(radius, radius), color=start_color_inhex)
|
|
bottom = Image.new('RGB', size=(radius, radius),color=end_color_inhex)
|
|
_canvas.paste(top, box=(0, 0, radius, radius))
|
|
_canvas.paste(g, box=(0, radius, radius, radius * 2))
|
|
_canvas.paste(bottom,box=(0, radius * 2, radius, radius * 3))
|
|
_canvas = _canvas.resize((radius * 3, radius * 3))
|
|
_canvas = __rotate_expand(_canvas,angle)
|
|
center = int(_canvas.width / 2)
|
|
_x = int(width / 3)
|
|
_y = int(height / 3)
|
|
ret_image = _canvas.crop((center - _x, center - _y, center + _x, center + _y))
|
|
ret_image = ret_image.resize((width, height))
|
|
return ret_image
|
|
|
|
def draw_rect(image:Image, x:int, y:int, width:int, height:int, line_color:str, line_width:int,
|
|
box_color:str=None) -> Image:
|
|
image = image.convert('RGBA')
|
|
draw = ImageDraw.Draw(image)
|
|
draw.rectangle((x, y, x + width, y + height), fill=box_color, outline=line_color, width=line_width, )
|
|
return image
|
|
|
|
|
|
'''Mask Functions'''
|
|
|
|
def expand_mask(mask:torch.Tensor, grow:int, blur:int) -> torch.Tensor:
|
|
# grow
|
|
c = 0
|
|
kernel = np.array([[c, 1, c],
|
|
[1, 1, 1],
|
|
[c, 1, c]])
|
|
growmask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
|
|
out = []
|
|
for m in growmask:
|
|
output = m.numpy()
|
|
for _ in range(abs(grow)):
|
|
if grow < 0:
|
|
output = scipy.ndimage.grey_erosion(output, footprint=kernel)
|
|
else:
|
|
output = scipy.ndimage.grey_dilation(output, footprint=kernel)
|
|
output = torch.from_numpy(output)
|
|
out.append(output)
|
|
# blur
|
|
for idx, tensor in enumerate(out):
|
|
pil_image = tensor2pil(tensor.cpu().detach())
|
|
pil_image = pil_image.filter(ImageFilter.GaussianBlur(blur))
|
|
out[idx] = pil2tensor(pil_image)
|
|
ret_mask = torch.cat(out, dim=0)
|
|
return ret_mask
|
|
|
|
def mask_invert(mask:torch.Tensor) -> torch.Tensor:
|
|
_image = mask2image(mask)
|
|
return image2mask(ImageChops.invert(_image))
|
|
|
|
def subtract_mask(masks_a:torch.Tensor, masks_b:torch.Tensor) -> torch.Tensor:
|
|
return torch.clamp(masks_a - masks_b, 0, 255)
|
|
|
|
def RGB2RGBA(image:Image, mask:Image) -> Image:
|
|
(R, G, B) = image.convert('RGB').split()
|
|
return Image.merge('RGBA', (R, G, B, mask.convert('L')))
|
|
|
|
def min_bounding_rect(image:Image) -> tuple:
|
|
cv2_image = pil2cv2(image)
|
|
gray = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2GRAY)
|
|
ret, thresh = cv2.threshold(gray, 127, 255, 0)
|
|
contours, _ = cv2.findContours(thresh, 1, 2)
|
|
x, y, width, height = 0, 0, 0, 0
|
|
area = 0
|
|
for contour in contours:
|
|
_x, _y, _w, _h = cv2.boundingRect(contour)
|
|
_area = _w * _h
|
|
if _area > area:
|
|
area = _area
|
|
x, y, width, height = _x, _y, _w, _h
|
|
|
|
return (x, y, width, height)
|
|
|
|
def max_inscribed_rect(image:Image) -> tuple:
|
|
img = pil2cv2(image)
|
|
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
ret, img_bin = cv2.threshold(img_gray, 127, 255, cv2.THRESH_BINARY)
|
|
contours, _ = cv2.findContours(img_bin, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
|
|
contour = contours[0].reshape(len(contours[0]), 2)
|
|
rect = []
|
|
for i in range(len(contour)):
|
|
x1, y1 = contour[i]
|
|
for j in range(len(contour)):
|
|
x2, y2 = contour[j]
|
|
area = abs(y2 - y1) * abs(x2 - x1)
|
|
rect.append(((x1, y1), (x2, y2), area))
|
|
all_rect = sorted(rect, key=lambda x: x[2], reverse=True)
|
|
if all_rect:
|
|
best_rect_found = False
|
|
index_rect = 0
|
|
nb_rect = len(all_rect)
|
|
while not best_rect_found and index_rect < nb_rect:
|
|
rect = all_rect[index_rect]
|
|
(x1, y1) = rect[0]
|
|
(x2, y2) = rect[1]
|
|
valid_rect = True
|
|
x = min(x1, x2)
|
|
while x < max(x1, x2) + 1 and valid_rect:
|
|
if any(img[y1, x]) == 0 or any(img[y2, x]) == 0:
|
|
valid_rect = False
|
|
x += 1
|
|
y = min(y1, y2)
|
|
while y < max(y1, y2) + 1 and valid_rect:
|
|
if any(img[y, x1]) == 0 or any(img[y, x2]) == 0:
|
|
valid_rect = False
|
|
y += 1
|
|
if valid_rect:
|
|
best_rect_found = True
|
|
index_rect += 1
|
|
#较小的数值排前面
|
|
log(f"x1={x1}, y1={y1},x2={x2}, y2={y2}")
|
|
x1, y1, x2, y2 = min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)
|
|
return (x1, y1, x2 - x1, y2 - y1)
|
|
|
|
'''Color Functions'''
|
|
|
|
def RGB_to_Hex(RGB) -> str:
|
|
color = '#'
|
|
for i in RGB:
|
|
num = int(i)
|
|
color += str(hex(num))[-2:].replace('x', '0').upper()
|
|
return color
|
|
|
|
def Hex_to_RGB(inhex) -> tuple:
|
|
rval = inhex[1:3]
|
|
gval = inhex[3:5]
|
|
bval = inhex[5:]
|
|
rgb = (int(rval, 16), int(gval, 16), int(bval, 16))
|
|
return tuple(rgb)
|
|
|
|
'''Value Functions'''
|
|
|
|
def step_value(start_value, end_value, total_step, step) -> float: # 按当前步数在总步数中的位置返回比例值
|
|
factor = step / total_step
|
|
return (end_value - start_value) * factor + start_value
|
|
|
|
def step_color(start_color_inhex:str, end_color_inhex:str, total_step:int, step:int) -> str: # 按当前步数在总步数中的位置返回比例颜色
|
|
start_color = tuple(Hex_to_RGB(start_color_inhex))
|
|
end_color = tuple(Hex_to_RGB(end_color_inhex))
|
|
start_R, start_G, start_B = start_color[0], start_color[1], start_color[2]
|
|
end_R, end_G, end_B = end_color[0], end_color[1], end_color[2]
|
|
ret_color = (int(step_value(start_R, end_R, total_step, step)),
|
|
int(step_value(start_G, end_G, total_step, step)),
|
|
int(step_value(start_B, end_B, total_step, step)),
|
|
)
|
|
return RGB_to_Hex(ret_color)
|
|
|
|
|