1866 lines
65 KiB
Python
1866 lines
65 KiB
Python
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import numpy as np
|
|
import torch.nn.functional as F
|
|
import cv2
|
|
from PIL import Image, ImageEnhance
|
|
from PIL import Image
|
|
from typing import Literal, Any
|
|
from torch import Tensor, dtype
|
|
import torchvision.transforms.functional as D
|
|
import math
|
|
|
|
def radialspace_2D(
|
|
size: tuple[int] | tuple[int, int],
|
|
curvy: float = 1.0,
|
|
scale: float = 1.0,
|
|
mode: str = "square",
|
|
min_val: float = 0.0,
|
|
max_val: float = 1.0,
|
|
center: tuple[float, float] = (0.5, 0.5),
|
|
function: Any = None,
|
|
normalize: bool = True,
|
|
dtype: dtype = torch.float32
|
|
) -> Tensor:
|
|
|
|
if isinstance(size, tuple):
|
|
if len(size) == 1:
|
|
width = height = size[0]
|
|
elif len(size) == 2:
|
|
width, height = size
|
|
else:
|
|
raise ValueError("Invalid size argument")
|
|
else:
|
|
raise TypeError("Size must be a tuple")
|
|
|
|
x = torch.linspace(0, 1, width)
|
|
y = torch.linspace(0, 1, height)
|
|
|
|
xx, yy = torch.meshgrid(x, y, indexing='ij')
|
|
|
|
if function is not None:
|
|
d = function(xx, yy)
|
|
elif mode == "square":
|
|
xx = (torch.abs(xx - center[0]) ** curvy)
|
|
yy = (torch.abs(yy - center[1]) ** curvy)
|
|
d = torch.max(xx, yy)
|
|
elif mode == "circle":
|
|
d = torch.sqrt((xx - center[0]) ** 2 + (yy - center[1]) ** 2)
|
|
d = (d ** curvy)
|
|
elif mode == "rectangle":
|
|
xx = (torch.abs(xx - center[0]) ** curvy)
|
|
yy = (torch.abs(yy - center[1]) ** curvy)
|
|
d = xx + yy
|
|
elif mode == "corners":
|
|
xx = (torch.abs(xx - center[0]) ** curvy)
|
|
yy = (torch.abs(yy - center[1]) ** curvy)
|
|
d = torch.min(xx, yy)
|
|
else:
|
|
raise ValueError("Not supported mode.")
|
|
|
|
if normalize:
|
|
d = d / d.max() * scale
|
|
d = torch.clamp(d, min_val, max_val)
|
|
|
|
return d.to(dtype)
|
|
|
|
def cv2_layer(tensor: Tensor, function) -> Tensor:
|
|
"""
|
|
This function applies a given function to each channel of an input tensor and returns the result as a PyTorch tensor.
|
|
|
|
:param tensor: A PyTorch tensor of shape (H, W, C) or (N, H, W, C), where C is the number of channels, H is the height, and W is the width of the image.
|
|
:param function: A function that takes a numpy array of shape (H, W, C) as input and returns a numpy array of the same shape.
|
|
:return: A PyTorch tensor of the same shape as the input tensor, where the given function has been applied to each channel of each image in the tensor.
|
|
"""
|
|
shape_size = tensor.shape.__len__()
|
|
|
|
def produce(image):
|
|
channels = image[0, 0, :].shape[0]
|
|
|
|
rgb = image[:, :, 0:3].numpy()
|
|
result_rgb = function(rgb)
|
|
|
|
if channels <= 3:
|
|
return torch.from_numpy(result_rgb)
|
|
elif channels == 4:
|
|
alpha = image[:, :, 3:4].numpy()
|
|
result_alpha = function(alpha)[..., np.newaxis]
|
|
result_rgba = np.concatenate((result_rgb, result_alpha), axis=2)
|
|
|
|
return torch.from_numpy(result_rgba)
|
|
|
|
if shape_size == 3:
|
|
return torch.from_numpy(produce(tensor))
|
|
elif shape_size == 4:
|
|
return torch.stack([
|
|
produce(tensor[i]) for i in range(len(tensor))
|
|
])
|
|
else:
|
|
raise ValueError("Incompatible tensor dimension.")
|
|
|
|
def apply_to_batch(func):
|
|
def wrapper(self, image, *args, **kwargs):
|
|
images = []
|
|
for img in image:
|
|
images.append(func(self, img, *args, **kwargs))
|
|
batch_tensor = torch.cat(images, dim=0)
|
|
return (batch_tensor, )
|
|
return wrapper
|
|
|
|
def gamma_correction_pil(image, gamma):
|
|
# Convert PIL Image to NumPy array
|
|
img_array = np.array(image)
|
|
# Normalization [0,255] -> [0,1]
|
|
img_array = img_array / 255.0
|
|
# Apply gamma correction
|
|
img_corrected = np.power(img_array, gamma)
|
|
# Convert corrected image back to original scale [0,1] -> [0,255]
|
|
img_corrected = np.uint8(img_corrected * 255)
|
|
# Convert NumPy array back to PIL Image
|
|
corrected_image = Image.fromarray(img_corrected)
|
|
return corrected_image
|
|
|
|
|
|
def apply_hald_clut(lut_image, img):
|
|
hald_w, hald_h = lut_image.size
|
|
clut_size = int(round(pow(hald_w, 1/3)))
|
|
scale = (clut_size * clut_size - 1) / 255
|
|
img = np.asarray(img)
|
|
|
|
# Convert the HaldCLUT image to numpy array
|
|
hald_img_array = np.asarray(lut_image)
|
|
|
|
# If the HaldCLUT image is monochrome, duplicate its single channel to three
|
|
if len(hald_img_array.shape) == 2:
|
|
hald_img_array = np.stack([hald_img_array]*3, axis=-1)
|
|
|
|
hald_img_array = hald_img_array.reshape(clut_size ** 6, 3)
|
|
|
|
clut_r = np.rint(img[:, :, 0] * scale).astype(int)
|
|
clut_g = np.rint(img[:, :, 1] * scale).astype(int)
|
|
clut_b = np.rint(img[:, :, 2] * scale).astype(int)
|
|
filtered_image = np.zeros((img.shape))
|
|
filtered_image[:, :] = hald_img_array[clut_r + clut_size ** 2 * clut_g + clut_size ** 4 * clut_b]
|
|
filtered_image = Image.fromarray(filtered_image.astype('uint8'), 'RGB')
|
|
return filtered_image
|
|
|
|
|
|
# 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)
|
|
|
|
def tensor2numpy(image):
|
|
# Convert tensor to numpy array and transpose dimensions from (C, H, W) to (H, W, C)
|
|
return (255.0 * image.cpu().numpy().squeeze().transpose(1, 2, 0)).astype(np.uint8)
|
|
|
|
|
|
|
|
class Blend:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
"blend_factor": ("FLOAT", {
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"blend_mode": ( [
|
|
"normal",
|
|
"multiply",
|
|
"screen",
|
|
"overlay",
|
|
"soft_light",
|
|
"hard_light",
|
|
"darken",
|
|
"lighten",
|
|
"color_dodge",
|
|
"color_burn",
|
|
"linear_dodge",
|
|
"linear_burn",
|
|
"linear_light",
|
|
"vivid_light",
|
|
"pin_light",
|
|
"difference",
|
|
"exclusion",
|
|
"divide"
|
|
],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "blend_images"
|
|
|
|
CATEGORY = "postprocessing/Blends"
|
|
|
|
def blend_images(
|
|
self,
|
|
image1: torch.Tensor,
|
|
image2: torch.Tensor,
|
|
blend_factor: float,
|
|
blend_mode: str,
|
|
):
|
|
if image1.shape != image2.shape:
|
|
image2 = self.crop_and_resize(image2, image1.shape)
|
|
|
|
blended_image = self.blend_mode(image1, image2, blend_mode)
|
|
blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor
|
|
blended_image = torch.clamp(blended_image, 0, 1)
|
|
return (blended_image,)
|
|
|
|
def blend_mode(self, img1, img2, mode):
|
|
if mode == "normal":
|
|
return img2
|
|
elif mode == "multiply":
|
|
return img1 * img2
|
|
elif mode == "divide":
|
|
return img1 / img2
|
|
elif mode == "screen":
|
|
return 1 - (1 - img1) * (1 - img2)
|
|
elif mode == "overlay":
|
|
return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
|
|
elif mode == "soft_light":
|
|
return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2-1)*(self.g(img1)-img1))
|
|
elif mode == 'hard_light':
|
|
return torch.where(img2 < 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
|
|
elif mode == 'darken':
|
|
return torch.min(img1, img2)
|
|
elif mode == 'lighten':
|
|
return torch.max(img1, img2)
|
|
elif mode == 'color_dodge':
|
|
return torch.where((img1 < 1), img1 / (1 - img2), torch.ones_like(img1))
|
|
elif mode == 'color_burn':
|
|
return torch.where(img2 > 0, (1 - ((1 - img1) / img2)), torch.zeros_like(img1))
|
|
elif mode == 'linear_dodge':
|
|
return img1 + img2
|
|
elif mode == 'linear_burn':
|
|
return img1 + img2 - 1
|
|
elif mode == 'linear_light':
|
|
return torch.where(img2 <= 0.5, img1 + 2 * (img2 - 0.5), img1 + 2 * img2 - 1)
|
|
elif mode == 'vivid_light':
|
|
return torch.where(img2 <= 0.5, (1 -(1-img1) / (2*img2)),img1/(2*(1-img2)))
|
|
elif mode == 'pin_light':
|
|
return torch.where(img2 < 0.5, torch.min(img1, 2*img2), torch.max(img1, 2*(img2-0.5)))
|
|
elif mode == 'difference':
|
|
return torch.abs(img1 - img2)
|
|
elif mode == 'exclusion':
|
|
return img1 + img2 - 2 * img1 * img2
|
|
else:
|
|
raise ValueError(f"Unsupported blend mode: {mode}")
|
|
|
|
def g(self, x):
|
|
return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x))
|
|
|
|
def crop_and_resize(self, img: torch.Tensor, target_shape: tuple):
|
|
batch_size, img_h, img_w, img_c = img.shape
|
|
_, target_h, target_w, _ = target_shape
|
|
img_aspect_ratio = img_w / img_h
|
|
target_aspect_ratio = target_w / target_h
|
|
|
|
# Crop center of the image to the target aspect ratio
|
|
if img_aspect_ratio > target_aspect_ratio:
|
|
new_width = int(img_h * target_aspect_ratio)
|
|
left = (img_w - new_width) // 2
|
|
img = img[:, :, left:left + new_width, :]
|
|
else:
|
|
new_height = int(img_w / target_aspect_ratio)
|
|
top = (img_h - new_height) // 2
|
|
img = img[:, top:top + new_height, :, :]
|
|
|
|
# Resize to target size
|
|
img = img.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
|
|
img = F.interpolate(img, size=(target_h, target_w), mode='bilinear', align_corners=False)
|
|
img = img.permute(0, 2, 3, 1)
|
|
|
|
return img
|
|
|
|
|
|
batch_size, img_h, img_w, img_c = img.shape
|
|
_, target_h, target_w, _ = target_shape
|
|
img_aspect_ratio = img_w / img_h
|
|
target_aspect_ratio = target_w / target_h
|
|
|
|
# Crop center of the image to the target aspect ratio
|
|
if img_aspect_ratio > target_aspect_ratio:
|
|
new_width = int(img_h * target_aspect_ratio)
|
|
left = (img_w - new_width) // 2
|
|
img = img[:, :, left:left + new_width, :]
|
|
else:
|
|
new_height = int(img_w / target_aspect_ratio)
|
|
top = (img_h - new_height) // 2
|
|
img = img[:, top:top + new_height, :, :]
|
|
|
|
# Resize to target size
|
|
img = img.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
|
|
img = F.interpolate(img, size=(target_h, target_w), mode='bilinear', align_corners=False)
|
|
img = img.permute(0, 2, 3, 1)
|
|
|
|
return img
|
|
|
|
|
|
class AlphaBlend:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
"image2_mask":("IMAGE",),
|
|
"blend_factor": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"blend_mode": ( [
|
|
"normal",
|
|
"add",
|
|
],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "blend_images"
|
|
|
|
CATEGORY = "postprocessing/Blends"
|
|
|
|
def blend_images(
|
|
self,
|
|
image1: torch.Tensor,
|
|
image2: torch.Tensor,
|
|
image2_mask: torch.Tensor,
|
|
blend_factor: float,
|
|
blend_mode: str,
|
|
):
|
|
|
|
image1 = image1.to('cuda:0')
|
|
image2 = image2.to('cuda:0')
|
|
image2_mask = image2_mask.to('cuda:0')
|
|
|
|
if image1.shape != image2.shape:
|
|
image2 = self.crop_and_resize(image2, image1.shape)
|
|
|
|
|
|
if image1.shape != image2_mask.shape:
|
|
image2_mask = self.crop_and_resize(image2_mask, image1.shape)
|
|
|
|
blended_image = self.blend_mode(image1, image2, image2_mask, blend_mode)
|
|
blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor
|
|
blended_image = torch.clamp(blended_image, 0, 1)
|
|
return (blended_image,)
|
|
|
|
def blend_mode(self, img1, img2,img2_mask, mode):
|
|
if mode == "normal":
|
|
return img1*img2_mask+img2*(1-img2_mask)
|
|
elif mode == "add":
|
|
return img1+img2*(1-img2_mask)
|
|
else:
|
|
raise ValueError(f"Unsupported blend mode: {mode}")
|
|
|
|
def g(self, x):
|
|
return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x))
|
|
|
|
def crop_and_resize(self, img: torch.Tensor, target_shape: tuple):
|
|
batch_size, img_h, img_w, img_c = img.shape
|
|
_, target_h, target_w, _ = target_shape
|
|
img_aspect_ratio = img_w / img_h
|
|
target_aspect_ratio = target_w / target_h
|
|
|
|
# Crop center of the image to the target aspect ratio
|
|
if img_aspect_ratio > target_aspect_ratio:
|
|
new_width = int(img_h * target_aspect_ratio)
|
|
left = (img_w - new_width) // 2
|
|
img = img[:, :, left:left + new_width, :]
|
|
else:
|
|
new_height = int(img_w / target_aspect_ratio)
|
|
top = (img_h - new_height) // 2
|
|
img = img[:, top:top + new_height, :, :]
|
|
|
|
# Resize to target size
|
|
img = img.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
|
|
img = F.interpolate(img, size=(target_h, target_w), mode='bilinear', align_corners=False)
|
|
img = img.permute(0, 2, 3, 1)
|
|
|
|
return img
|
|
|
|
|
|
batch_size, img_h, img_w, img_c = img.shape
|
|
_, target_h, target_w, _ = target_shape
|
|
img_aspect_ratio = img_w / img_h
|
|
target_aspect_ratio = target_w / target_h
|
|
|
|
# Crop center of the image to the target aspect ratio
|
|
if img_aspect_ratio > target_aspect_ratio:
|
|
new_width = int(img_h * target_aspect_ratio)
|
|
left = (img_w - new_width) // 2
|
|
img = img[:, :, left:left + new_width, :]
|
|
else:
|
|
new_height = int(img_w / target_aspect_ratio)
|
|
top = (img_h - new_height) // 2
|
|
img = img[:, top:top + new_height, :, :]
|
|
|
|
# Resize to target size
|
|
img = img.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
|
|
img = F.interpolate(img, size=(target_h, target_w), mode='bilinear', align_corners=False)
|
|
img = img.permute(0, 2, 3, 1)
|
|
|
|
return img
|
|
|
|
class Blur:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"blur_radius": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 30,
|
|
"step": 1
|
|
}),
|
|
"sigma": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.1,
|
|
"max": 10.0,
|
|
"step": 0.1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "blur"
|
|
|
|
CATEGORY = "postprocessing/Filters"
|
|
|
|
def blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
|
|
if blur_radius == 0:
|
|
return (image,)
|
|
|
|
batch_size, height, width, channels = image.shape
|
|
|
|
kernel_size = blur_radius * 2 + 1
|
|
kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
|
|
|
|
image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
|
|
blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels)
|
|
blurred = blurred.permute(0, 2, 3, 1)
|
|
|
|
return (blurred,)
|
|
|
|
class CannyEdgeMask:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"lower_threshold": ("FLOAT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 1,
|
|
"step": 0.01
|
|
}),
|
|
"upper_threshold": ("FLOAT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 1,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "canny"
|
|
|
|
CATEGORY = "postprocessing/Masks"
|
|
|
|
def canny(self, image: torch.Tensor, lower_threshold: float, upper_threshold: float):
|
|
batch_size, height, width, _ = image.shape
|
|
result = torch.zeros(batch_size, height, width)
|
|
lower_threshold*=500
|
|
upper_threshold*=500
|
|
|
|
for b in range(batch_size):
|
|
tensor_image = image[b].numpy().copy()
|
|
gray_image = (cv2.cvtColor(tensor_image, cv2.COLOR_RGB2GRAY) * 255).astype(np.uint8)
|
|
canny = cv2.Canny(gray_image, lower_threshold, upper_threshold)
|
|
tensor = torch.from_numpy(canny)
|
|
result[b] = tensor
|
|
|
|
return (result,)
|
|
|
|
class ChromaticAberration:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"red_shift": ("INT", {
|
|
"default": 0,
|
|
"min": -20,
|
|
"max": 20,
|
|
"step": 1
|
|
}),
|
|
"red_direction": (["horizontal", "vertical"],),
|
|
"green_shift": ("INT", {
|
|
"default": 0,
|
|
"min": -20,
|
|
"max": 20,
|
|
"step": 1
|
|
}),
|
|
"green_direction": (["horizontal", "vertical"],),
|
|
"blue_shift": ("INT", {
|
|
"default": 0,
|
|
"min": -20,
|
|
"max": 20,
|
|
"step": 1
|
|
}),
|
|
"blue_direction": (["horizontal", "vertical"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "chromatic_aberration"
|
|
|
|
CATEGORY = "postprocessing/Effects"
|
|
|
|
def chromatic_aberration(self, image: torch.Tensor, red_shift: int, green_shift: int, blue_shift: int, red_direction: str, green_direction: str, blue_direction: str):
|
|
def get_shift(direction, shift):
|
|
shift = -shift if direction == 'vertical' else shift # invert vertical shift as otherwise positive actually shifts down
|
|
return (shift, 0) if direction == 'vertical' else (0, shift)
|
|
|
|
x = image.permute(0, 3, 1, 2)
|
|
shifts = [get_shift(direction, shift) for direction, shift in zip([red_direction, green_direction, blue_direction], [red_shift, green_shift, blue_shift])]
|
|
channels = [torch.roll(x[:, i, :, :], shifts=shifts[i], dims=(1, 2)) for i in range(3)]
|
|
|
|
output = torch.stack(channels, dim=1)
|
|
output = output.permute(0, 2, 3, 1)
|
|
|
|
return (output,)
|
|
|
|
#https://github.com/hahnec/color-matcher/
|
|
from color_matcher import ColorMatcher
|
|
|
|
class ColorMatch:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image_ref": ("IMAGE",),
|
|
"image_target": ("IMAGE",),
|
|
"method": (
|
|
[
|
|
'mkl',
|
|
'hm',
|
|
'reinhard',
|
|
'mvgd',
|
|
'hm-mvgd-hm',
|
|
'hm-mkl-hm',
|
|
], {
|
|
"default": 'mkl'
|
|
}),
|
|
|
|
},
|
|
}
|
|
|
|
CATEGORY = "postprocessing/Color Adjustments"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "colormatch"
|
|
|
|
def colormatch(self, image_ref, image_target, method):
|
|
cm = ColorMatcher()
|
|
image_ref = image_ref.cpu()
|
|
image_target = image_target.cpu()
|
|
batch_size = image_target.size(0)
|
|
out = []
|
|
images_target = image_target.squeeze()
|
|
images_ref = image_ref.squeeze()
|
|
|
|
image_ref_np = images_ref.numpy()
|
|
images_target_np = images_target.numpy()
|
|
|
|
if image_ref.size(0) > 1 and image_ref.size(0) != batch_size:
|
|
raise ValueError("ColorMatch: Use either single reference image or a matching batch of reference images.")
|
|
|
|
for i in range(batch_size):
|
|
image_target_np = images_target_np if batch_size == 1 else images_target[i].numpy()
|
|
image_ref_np_i = image_ref_np if image_ref.size(0) == 1 else images_ref[i].numpy()
|
|
try:
|
|
image_result = cm.transfer(src=image_target_np, ref=image_ref_np_i, method=method)
|
|
except BaseException as e:
|
|
print(f"Error occurred during transfer: {e}")
|
|
break
|
|
out.append(torch.from_numpy(image_result))
|
|
return (torch.stack(out, dim=0).to(torch.float32), )
|
|
|
|
class ColorCorrect:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"temperature": ("FLOAT", {
|
|
"default": 0,
|
|
"min": -1,
|
|
"max": 1,
|
|
"step": 0.01
|
|
}),
|
|
"hue": ("FLOAT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 360,
|
|
"step": 1
|
|
}),
|
|
"brightness": ("FLOAT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 5,
|
|
"step": 0.01
|
|
}),
|
|
"contrast": ("FLOAT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 5,
|
|
"step": 0.01
|
|
}),
|
|
"saturation": ("FLOAT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 5,
|
|
"step": 0.01
|
|
}),
|
|
"gamma": ("FLOAT", {
|
|
"default": 1,
|
|
"min": 0.2,
|
|
"max": 2.2,
|
|
"step": 0.1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "color_correct"
|
|
|
|
CATEGORY = "postprocessing/Color Adjustments"
|
|
|
|
def color_correct(self, image: torch.Tensor, temperature: float, hue: float, brightness: float, contrast: float, saturation: float, gamma: float):
|
|
batch_size, height, width, _ = image.shape
|
|
result = torch.zeros_like(image)
|
|
|
|
for b in range(batch_size):
|
|
tensor_image = image[b].numpy()
|
|
|
|
modified_image = Image.fromarray((tensor_image * 255).astype(np.uint8))
|
|
|
|
# brightness
|
|
modified_image = ImageEnhance.Brightness(modified_image).enhance(brightness)
|
|
|
|
# contrast
|
|
modified_image = ImageEnhance.Contrast(modified_image).enhance(contrast)
|
|
modified_image = np.array(modified_image).astype(np.float32)
|
|
|
|
# temperature
|
|
if temperature > 0:
|
|
modified_image[:, :, 0] *= 1 + temperature
|
|
modified_image[:, :, 1] *= 1 + temperature * 0.4
|
|
elif temperature < 0:
|
|
modified_image[:, :, 2] *= 1 - temperature
|
|
modified_image = np.clip(modified_image, 0, 255)/255
|
|
|
|
# gamma
|
|
modified_image = np.clip(np.power(modified_image, gamma), 0, 1)
|
|
|
|
# saturation
|
|
hls_img = cv2.cvtColor(modified_image, cv2.COLOR_RGB2HLS)
|
|
hls_img[:, :, 2] = np.clip(saturation*hls_img[:, :, 2], 0, 1)
|
|
modified_image = cv2.cvtColor(hls_img, cv2.COLOR_HLS2RGB) * 255
|
|
|
|
# hue
|
|
hsv_img = cv2.cvtColor(modified_image, cv2.COLOR_RGB2HSV)
|
|
hsv_img[:, :, 0] = (hsv_img[:, :, 0] + hue) % 360
|
|
modified_image = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2RGB)
|
|
|
|
modified_image = modified_image.astype(np.uint8)
|
|
modified_image = modified_image / 255
|
|
modified_image = torch.from_numpy(modified_image).unsqueeze(0)
|
|
result[b] = modified_image
|
|
|
|
return (result, )
|
|
|
|
class ColorTint:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.1
|
|
}),
|
|
"mode": (["custom","sepia", "red", "green", "blue", "cyan", "magenta", "yellow", "purple", "orange", "warm", "cool", "lime", "navy", "vintage", "rose", "teal", "maroon", "peach", "lavender", "olive"],),
|
|
"tint_color_hex": ("STRING", {"forceInput": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "color_tint"
|
|
|
|
CATEGORY = "postprocessing/Color Adjustments"
|
|
|
|
# def hex_to_rgb(hex_color):
|
|
# hex_color = hex_color.lstrip('#') # Remove the '#' character, if present
|
|
# r = int(hex_color[0:2], 16)
|
|
# g = int(hex_color[2:4], 16)
|
|
# b = int(hex_color[4:6], 16)
|
|
# return (r, g, b)
|
|
|
|
def color_tint(self, image: torch.Tensor, strength: float, mode: str = "sepia" , tint_color_hex='#000000'):
|
|
if strength == 0:
|
|
return (image,)
|
|
|
|
|
|
tint_color_hex = tint_color_hex.lstrip('#')
|
|
|
|
|
|
color_rgb = [int(tint_color_hex[0:2], 16)/255,int(tint_color_hex[2:4], 16)/255,int(tint_color_hex[4:6], 16)/255]
|
|
|
|
|
|
|
|
sepia_weights = torch.tensor([0.2989, 0.5870, 0.1140]).view(1, 1, 1, 3).to(image.device)
|
|
|
|
mode_filters = {
|
|
"custom": torch.tensor([color_rgb[0], color_rgb[1], color_rgb[2]]),
|
|
"sepia": torch.tensor([1.0, 0.8, 0.6]),
|
|
"red": torch.tensor([1.0, 0.6, 0.6]),
|
|
"green": torch.tensor([0.6, 1.0, 0.6]),
|
|
"blue": torch.tensor([0.6, 0.8, 1.0]),
|
|
"cyan": torch.tensor([0.6, 1.0, 1.0]),
|
|
"magenta": torch.tensor([1.0, 0.6, 1.0]),
|
|
"yellow": torch.tensor([1.0, 1.0, 0.6]),
|
|
"purple": torch.tensor([0.8, 0.6, 1.0]),
|
|
"orange": torch.tensor([1.0, 0.7, 0.3]),
|
|
"warm": torch.tensor([1.0, 0.9, 0.7]),
|
|
"cool": torch.tensor([0.7, 0.9, 1.0]),
|
|
"lime": torch.tensor([0.7, 1.0, 0.3]),
|
|
"navy": torch.tensor([0.3, 0.4, 0.7]),
|
|
"vintage": torch.tensor([0.9, 0.85, 0.7]),
|
|
"rose": torch.tensor([1.0, 0.8, 0.9]),
|
|
"teal": torch.tensor([0.3, 0.8, 0.8]),
|
|
"maroon": torch.tensor([0.7, 0.3, 0.5]),
|
|
"peach": torch.tensor([1.0, 0.8, 0.6]),
|
|
"lavender": torch.tensor([0.8, 0.6, 1.0]),
|
|
"olive": torch.tensor([0.6, 0.7, 0.4]),
|
|
}
|
|
|
|
scale_filter = mode_filters[mode].view(1, 1, 1, 3).to(image.device)
|
|
|
|
grayscale = torch.sum(image * sepia_weights, dim=-1, keepdim=True)
|
|
tinted = grayscale * scale_filter
|
|
|
|
result = tinted * strength + image * (1 - strength)
|
|
return (result,)
|
|
|
|
class FilmGrain:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"intensity": ("FLOAT", {
|
|
"default": 0.2,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"scale": ("FLOAT", {
|
|
"default": 10,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1
|
|
}),
|
|
"temperature": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": -100,
|
|
"max": 100,
|
|
"step": 1
|
|
}),
|
|
"vignette": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 1.0
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "film_grain"
|
|
|
|
CATEGORY = "postprocessing/Effects"
|
|
|
|
def film_grain(self, image: torch.Tensor, intensity: float, scale: float, temperature: float, vignette: float):
|
|
batch_size, height, width, _ = image.shape
|
|
result = torch.zeros_like(image)
|
|
|
|
for b in range(batch_size):
|
|
tensor_image = image[b].numpy()
|
|
|
|
# Generate Perlin noise with shape (height, width) and scale
|
|
noise = self.generate_perlin_noise((height, width), scale)
|
|
noise = (noise - np.min(noise)) / (np.max(noise) - np.min(noise))
|
|
|
|
# Apply grain intensity
|
|
noise = (noise * 2 - 1) * intensity
|
|
|
|
# Blend the noise with the image
|
|
grain_image = np.clip(tensor_image + noise[:, :, np.newaxis], 0, 1)
|
|
|
|
# Apply temperature
|
|
grain_image = self.apply_temperature(grain_image, temperature)
|
|
|
|
# Apply vignette
|
|
grain_image = self.apply_vignette(grain_image, vignette)
|
|
|
|
tensor = torch.from_numpy(grain_image).unsqueeze(0)
|
|
result[b] = tensor
|
|
|
|
return (result,)
|
|
|
|
def generate_perlin_noise(self, shape, scale, octaves=4, persistence=0.5, lacunarity=2):
|
|
def smoothstep(t):
|
|
return t * t * (3.0 - 2.0 * t)
|
|
|
|
def lerp(t, a, b):
|
|
return a + t * (b - a)
|
|
|
|
def gradient(h, x, y):
|
|
vectors = np.array([[1, 1], [-1, 1], [1, -1], [-1, -1]])
|
|
g = vectors[h % 4]
|
|
return g[:, :, 0] * x + g[:, :, 1] * y
|
|
|
|
height, width = shape
|
|
noise = np.zeros(shape)
|
|
|
|
for octave in range(octaves):
|
|
octave_scale = scale * lacunarity ** octave
|
|
x = np.linspace(0, 1, width, endpoint=False)
|
|
y = np.linspace(0, 1, height, endpoint=False)
|
|
X, Y = np.meshgrid(x, y)
|
|
X, Y = X * octave_scale, Y * octave_scale
|
|
|
|
xi = X.astype(int)
|
|
yi = Y.astype(int)
|
|
|
|
xf = X - xi
|
|
yf = Y - yi
|
|
|
|
u = smoothstep(xf)
|
|
v = smoothstep(yf)
|
|
|
|
n00 = gradient(np.random.randint(0, 4, (height, width)), xf, yf)
|
|
n01 = gradient(np.random.randint(0, 4, (height, width)), xf, yf - 1)
|
|
n10 = gradient(np.random.randint(0, 4, (height, width)), xf - 1, yf)
|
|
n11 = gradient(np.random.randint(0, 4, (height, width)), xf - 1, yf - 1)
|
|
|
|
x1 = lerp(u, n00, n10)
|
|
x2 = lerp(u, n01, n11)
|
|
y1 = lerp(v, x1, x2)
|
|
|
|
noise += y1 * persistence ** octave
|
|
|
|
return noise / (1 - persistence ** octaves)
|
|
|
|
def apply_temperature(self, image, temperature):
|
|
if temperature == 0:
|
|
return image
|
|
|
|
temperature /= 100
|
|
|
|
new_image = image.copy()
|
|
|
|
if temperature > 0:
|
|
new_image[:, :, 0] *= 1 + temperature
|
|
new_image[:, :, 1] *= 1 + temperature * 0.4
|
|
else:
|
|
new_image[:, :, 2] *= 1 - temperature
|
|
|
|
return np.clip(new_image, 0, 1)
|
|
|
|
def apply_vignette(self, image, vignette_strength):
|
|
if vignette_strength == 0:
|
|
return image
|
|
|
|
height, width, _ = image.shape
|
|
x = np.linspace(-1, 1, width)
|
|
y = np.linspace(-1, 1, height)
|
|
X, Y = np.meshgrid(x, y)
|
|
radius = np.sqrt(X ** 2 + Y ** 2)
|
|
|
|
# Map vignette strength from 0-10 to 1.800-0.800
|
|
mapped_vignette_strength = 1.8 - (vignette_strength - 1) * 0.1
|
|
vignette = 1 - np.clip(radius / mapped_vignette_strength, 0, 1)
|
|
|
|
return np.clip(image * vignette[..., np.newaxis], 0, 1)
|
|
|
|
class Glow:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"intensity": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 5.0,
|
|
"step": 0.01
|
|
}),
|
|
"blur_radius": ("INT", {
|
|
"default": 5,
|
|
"min": 1,
|
|
"max": 50,
|
|
"step": 1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply_glow"
|
|
|
|
CATEGORY = "postprocessing/Effects"
|
|
|
|
def apply_glow(self, image: torch.Tensor, intensity: float, blur_radius: int):
|
|
blurred_image = self.gaussian_blur(image, 2 * blur_radius + 1)
|
|
glowing_image = self.add_glow(image, blurred_image, intensity)
|
|
glowing_image = torch.clamp(glowing_image, 0, 1)
|
|
return (glowing_image,)
|
|
|
|
def gaussian_blur(self, image: torch.Tensor, kernel_size: int):
|
|
batch_size, height, width, channels = image.shape
|
|
|
|
sigma = (kernel_size - 1) / 6
|
|
kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
|
|
|
|
image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
|
|
blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels)
|
|
blurred = blurred.permute(0, 2, 3, 1)
|
|
|
|
return blurred
|
|
|
|
def add_glow(self, img, blurred_img, intensity):
|
|
return img + blurred_img * intensity
|
|
|
|
class HSVThresholdMask:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"low_threshold": ("FLOAT", {
|
|
"default": 0.2,
|
|
"min": 0,
|
|
"max": 1,
|
|
"step": 0.1
|
|
}),
|
|
"high_threshold": ("FLOAT", {
|
|
"default": 0.7,
|
|
"min": 0,
|
|
"max": 1,
|
|
"step": 0.1
|
|
}),
|
|
"hsv_channel": (["hue", "saturation", "value"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "hsv_threshold"
|
|
|
|
CATEGORY = "postprocessing/Masks"
|
|
|
|
def hsv_threshold(self, image: torch.Tensor, low_threshold: float, high_threshold: float, hsv_channel: str):
|
|
batch_size, height, width, _ = image.shape
|
|
result = torch.zeros(batch_size, height, width)
|
|
|
|
if hsv_channel == "hue":
|
|
channel = 0
|
|
low_threshold, high_threshold = int(low_threshold * 180), int(high_threshold * 180)
|
|
elif hsv_channel == "saturation":
|
|
channel = 1
|
|
low_threshold, high_threshold = int(low_threshold * 255), int(high_threshold * 255)
|
|
elif hsv_channel == "value":
|
|
channel = 2
|
|
low_threshold, high_threshold = int(low_threshold * 255), int(high_threshold * 255)
|
|
|
|
for b in range(batch_size):
|
|
tensor_image = (image[b].numpy().copy() * 255).astype(np.uint8)
|
|
hsv_image = cv2.cvtColor(tensor_image, cv2.COLOR_RGB2HSV)
|
|
|
|
mask = cv2.inRange(hsv_image[:, :, channel], low_threshold, high_threshold)
|
|
tensor = torch.from_numpy(mask).float() / 255.
|
|
result[b] = tensor
|
|
|
|
return (result,)
|
|
|
|
class KuwaharaBlur:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"blur_radius": ("INT", {
|
|
"default": 3,
|
|
"min": 0,
|
|
"max": 31,
|
|
"step": 1
|
|
}),
|
|
"method": (["mean", "gaussian"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply_kuwahara_filter"
|
|
|
|
CATEGORY = "postprocessing/Filters"
|
|
|
|
def apply_kuwahara_filter(self, image: np.ndarray, blur_radius: int, method: str):
|
|
if blur_radius == 0:
|
|
return (image,)
|
|
|
|
out = torch.zeros_like(image)
|
|
batch_size, height, width, channels = image.shape
|
|
|
|
for b in range(batch_size):
|
|
image = image[b].cpu().numpy() * 255.0
|
|
image = image.astype(np.uint8)
|
|
|
|
out[b] = torch.from_numpy(kuwahara(image, method=method, radius=blur_radius)) / 255.0
|
|
|
|
return (out,)
|
|
|
|
def kuwahara(orig_img, method="mean", radius=3, sigma=None):
|
|
if method == "gaussian" and sigma is None:
|
|
sigma = -1
|
|
|
|
image = orig_img.astype(np.float32, copy=False)
|
|
avgs = np.empty((4, *image.shape), dtype=image.dtype)
|
|
stddevs = np.empty((4, *image.shape[:2]), dtype=image.dtype)
|
|
image_2d = cv2.cvtColor(orig_img, cv2.COLOR_BGR2GRAY).astype(image.dtype, copy=False)
|
|
avgs_2d = np.empty((4, *image.shape[:2]), dtype=image.dtype)
|
|
|
|
squared_img = image_2d ** 2
|
|
|
|
if method == "mean":
|
|
kxy = np.ones(radius + 1, dtype=image.dtype) / (radius + 1)
|
|
elif method == "gaussian":
|
|
kxy = cv2.getGaussianKernel(2 * radius + 1, sigma, ktype=cv2.CV_32F)
|
|
kxy /= kxy[radius:].sum()
|
|
klr = np.array([kxy[:radius+1], kxy[radius:]])
|
|
kindexes = [[1, 1], [1, 0], [0, 1], [0, 0]]
|
|
|
|
shift = [(0, 0), (0, radius), (radius, 0), (radius, radius)]
|
|
|
|
for k in range(4):
|
|
if method == "mean":
|
|
kx, ky = kxy, kxy
|
|
else:
|
|
kx, ky = klr[kindexes[k]]
|
|
cv2.sepFilter2D(image, -1, kx, ky, avgs[k], shift[k])
|
|
cv2.sepFilter2D(image_2d, -1, kx, ky, avgs_2d[k], shift[k])
|
|
cv2.sepFilter2D(squared_img, -1, kx, ky, stddevs[k], shift[k])
|
|
stddevs[k] = stddevs[k] - avgs_2d[k] ** 2
|
|
|
|
indices = np.argmin(stddevs, axis=0)
|
|
filtered = np.take_along_axis(avgs, indices[None,...,None], 0).reshape(image.shape)
|
|
|
|
return filtered.astype(orig_img.dtype)
|
|
|
|
|
|
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mask": ("IMAGE",),
|
|
"direction": (["horizontal", "vertical"],),
|
|
"span_limit": ("INT", {
|
|
"default": None,
|
|
"min": 0,
|
|
"max": 100,
|
|
"step": 5
|
|
}),
|
|
"sort_by": (["hue", "saturation", "value"],),
|
|
"order": (["forward", "backward"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "sort_pixels"
|
|
|
|
CATEGORY = "postprocessing/Effects"
|
|
|
|
def sort_pixels(self, image: torch.Tensor, mask: torch.Tensor, direction: str, span_limit: int, sort_by: str, order: str):
|
|
horizontal_sort = direction == "horizontal"
|
|
reverse_sorting = order == "backward"
|
|
sort_by = sort_by[0].upper()
|
|
span_limit = span_limit if span_limit > 0 else None
|
|
|
|
batch_size = image.shape[0]
|
|
result = torch.zeros_like(image)
|
|
|
|
for b in range(batch_size):
|
|
tensor_img = image[b].numpy()
|
|
tensor_mask = mask[b].numpy()
|
|
sorted_image = pixel_sort(tensor_img, tensor_mask, horizontal_sort, span_limit, sort_by, reverse_sorting)
|
|
result[b] = torch.from_numpy(sorted_image)
|
|
|
|
return (result,)
|
|
|
|
class Pixelize:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"pixel_size": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 256,
|
|
"step": 1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply_pixelize"
|
|
|
|
CATEGORY = "postprocessing/Effects"
|
|
|
|
def apply_pixelize(self, image: torch.Tensor, pixel_size: int):
|
|
pixelized_image = self.pixelize_image(image, pixel_size)
|
|
pixelized_image = torch.clamp(pixelized_image, 0, 1)
|
|
return (pixelized_image,)
|
|
|
|
def pixelize_image(self, image: torch.Tensor, pixel_size: int):
|
|
batch_size, height, width, channels = image.shape
|
|
new_height = height // pixel_size
|
|
new_width = width // pixel_size
|
|
|
|
image = image.permute(0, 3, 1, 2)
|
|
image = F.avg_pool2d(image, kernel_size=pixel_size, stride=pixel_size)
|
|
image = F.interpolate(image, size=(height, width), mode='nearest')
|
|
image = image.permute(0, 2, 3, 1)
|
|
|
|
return image
|
|
|
|
class Quantize:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"colors": ("INT", {
|
|
"default": 16,
|
|
"min": 1,
|
|
"max": 256,
|
|
"step": 1
|
|
}),
|
|
"dither": (["none", "floyd-steinberg"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "quantize"
|
|
|
|
CATEGORY = "postprocessing/Color Adjustments"
|
|
|
|
def quantize(self, image: torch.Tensor, colors: int = 256, dither: str = "FLOYDSTEINBERG"):
|
|
batch_size, height, width, _ = image.shape
|
|
result = torch.zeros_like(image)
|
|
|
|
dither_option = Image.Dither.FLOYDSTEINBERG if dither == "floyd-steinberg" else Image.Dither.NONE
|
|
|
|
for b in range(batch_size):
|
|
tensor_image = image[b]
|
|
img = (tensor_image * 255).to(torch.uint8).numpy()
|
|
pil_image = Image.fromarray(img, mode='RGB')
|
|
|
|
palette = pil_image.quantize(colors=colors) # Required as described in https://github.com/python-pillow/Pillow/issues/5836
|
|
quantized_image = pil_image.quantize(colors=colors, palette=palette, dither=dither_option)
|
|
|
|
quantized_array = torch.tensor(np.array(quantized_image.convert("RGB"))).float() / 255
|
|
result[b] = quantized_array
|
|
|
|
return (result,)
|
|
|
|
|
|
class SineWave:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"amplitude": ("FLOAT", {
|
|
"default": 20,
|
|
"min": 0,
|
|
"max": 150,
|
|
"step": 5
|
|
}),
|
|
"frequency": ("FLOAT", {
|
|
"default": 10,
|
|
"min": 0,
|
|
"max": 20,
|
|
"step": 1
|
|
}),
|
|
"direction": (["horizontal", "vertical"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply_sine_wave"
|
|
|
|
CATEGORY = "postprocessing/Effects"
|
|
|
|
def apply_sine_wave(self, image: torch.Tensor, amplitude: float, frequency: float, direction: str):
|
|
batch_size, height, width, channels = image.shape
|
|
result = torch.zeros_like(image)
|
|
|
|
for b in range(batch_size):
|
|
tensor_image = image[b]
|
|
result[b] = self.sine_wave_effect(tensor_image, amplitude, frequency, direction)
|
|
|
|
return (result,)
|
|
|
|
def sine_wave_effect(self, image: torch.Tensor, amplitude: float, frequency: float, direction: str):
|
|
height, width, _ = image.shape
|
|
shifted_image = torch.zeros_like(image)
|
|
|
|
for channel in range(3):
|
|
if direction == "horizontal":
|
|
for i in range(height):
|
|
offset = int(amplitude * np.sin(2 * torch.pi * i * frequency / height))
|
|
shifted_image[i, :, channel] = torch.roll(image[i, :, channel], offset)
|
|
elif direction == "vertical":
|
|
for j in range(width):
|
|
offset = int(amplitude * np.sin(2 * torch.pi * j * frequency / width))
|
|
shifted_image[:, j, channel] = torch.roll(image[:, j, channel], offset)
|
|
|
|
return shifted_image
|
|
|
|
class Solarize:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"threshold": ("FLOAT", {
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "solarize_image"
|
|
|
|
CATEGORY = "postprocessing/Color Adjustments"
|
|
|
|
def solarize_image(self, image: torch.Tensor, threshold: float):
|
|
solarized_image = torch.where(image > threshold, 1 - image, image)
|
|
solarized_image = torch.clamp(solarized_image, 0, 1)
|
|
return (solarized_image,)
|
|
|
|
class LensBokeh:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"blades_shape": ("INT", {
|
|
"default": 5,
|
|
"min": 3,
|
|
}),
|
|
"blades_radius": ("INT", {
|
|
"default": 10,
|
|
"min": 1,
|
|
}),
|
|
"blades_rotation": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 360.0,
|
|
}),
|
|
"blur_size": ("INT", {
|
|
"default": 10,
|
|
"min": 1,
|
|
"step": 2
|
|
}),
|
|
"blur_type": (["bilateral", "stack", "none"],),
|
|
"method": (["dilate", "filter"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "postprocessing/lens"
|
|
|
|
# noinspection PyUnresolvedReferences
|
|
def lens_blur(self, image, blades_shape, blades_radius, blades_rotation, method):
|
|
angles = np.linspace(0, 2 * np.pi, blades_shape + 1)[:-1] + blades_rotation * np.pi / 180
|
|
x = blades_radius * np.cos(angles) + blades_radius
|
|
y = blades_radius * np.sin(angles) + blades_radius
|
|
pts = np.stack([x, y], axis=1).astype(np.int32)
|
|
|
|
mask = np.zeros((blades_radius * 2 + 1, blades_radius * 2 + 1), np.uint8)
|
|
cv2.fillPoly(mask, [pts], 255)
|
|
|
|
gaussian_kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
|
|
|
|
if method == "dilate":
|
|
kernel = cv2.filter2D(mask, -1, gaussian_kernel)
|
|
result = cv2.dilate(image, kernel)
|
|
elif method == "filter":
|
|
height, width = image.shape[:2]
|
|
dilate_size = min(height, width) // 512
|
|
|
|
if dilate_size > 0:
|
|
image = cv2.dilate(image, np.ones((dilate_size, dilate_size), np.uint8))
|
|
|
|
kernel = mask.astype(np.float32) / np.sum(mask)
|
|
kernel = cv2.filter2D(kernel, -1, gaussian_kernel)
|
|
result = cv2.filter2D(image, -1, kernel)
|
|
else:
|
|
raise ValueError("Unsupported method.")
|
|
|
|
return result
|
|
|
|
def node(self, images, blades_shape, blades_radius, blades_rotation, blur_size, blur_type, method):
|
|
tensor = images.clone().detach()
|
|
blur_size -= 1
|
|
|
|
if blur_type == "bilateral":
|
|
tensor = cv2_layer(tensor, lambda x: cv2.bilateralFilter(x, blur_size, -100, 100))
|
|
elif blur_type == "stack":
|
|
tensor = cv2_layer(tensor, lambda x: cv2.stackBlur(x, (blur_size, blur_size)))
|
|
elif blur_type == "none":
|
|
pass
|
|
else:
|
|
raise ValueError("Unsupported blur type.")
|
|
|
|
return (cv2_layer(tensor, lambda x: self.lens_blur(
|
|
x, blades_shape, blades_radius, blades_rotation, method
|
|
)),)
|
|
|
|
class LensOpticAxis:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"lens_shape": (["circle", "square", "rectangle", "corners"],),
|
|
"lens_edge": (["around", "symmetric"],),
|
|
"lens_curvy": ("FLOAT", {
|
|
"default": 4.0,
|
|
"max": 15.0,
|
|
"step": 0.1,
|
|
}),
|
|
"lens_zoom": ("FLOAT", {
|
|
"default": 2.0,
|
|
"step": 0.1,
|
|
}),
|
|
"lens_aperture": ("FLOAT", {
|
|
"default": 0.5,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
}),
|
|
"blur_intensity": ("INT", {
|
|
"default": 30,
|
|
"min": 2,
|
|
"step": 2
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "node"
|
|
CATEGORY = "postprocessing/lens"
|
|
|
|
def node(self, images, lens_shape, lens_edge, lens_curvy, lens_zoom, lens_aperture, blur_intensity):
|
|
blur_intensity -= 1
|
|
lens_zoom += 1
|
|
|
|
height, width = images[0, :, :, 0].shape
|
|
|
|
if lens_edge == "around":
|
|
mask = radialspace_2D((height, width), lens_curvy, lens_zoom, lens_shape, 0.0, 1.0 + lens_curvy).unsqueeze(0).unsqueeze(3)
|
|
elif lens_edge == "symmetric":
|
|
if height != width:
|
|
new_height = new_width = max(height, width)
|
|
crop_top_bottom = (new_height - height) // 2
|
|
crop_left_right = (new_width - width) // 2
|
|
|
|
mask = radialspace_2D((new_height, new_width), lens_curvy, lens_zoom, lens_shape, 0.0, 1.0 + lens_curvy)[
|
|
crop_top_bottom:-crop_top_bottom or None,
|
|
crop_left_right:-crop_left_right or None
|
|
].unsqueeze(0).unsqueeze(3)
|
|
else:
|
|
mask = radialspace_2D((height, width), lens_curvy, lens_zoom, lens_shape, 0.0, 1.0 + lens_curvy).unsqueeze(0).unsqueeze(3)
|
|
else:
|
|
raise ValueError("Not existing lens_edge parameter.")
|
|
|
|
center_x = width // 2
|
|
center_y = height // 2
|
|
|
|
y_v, x_v = torch.meshgrid(torch.arange(height), torch.arange(width), indexing='ij')
|
|
|
|
dx = x_v - center_x
|
|
dy = y_v - center_y
|
|
|
|
distance = torch.sqrt(dx ** 2 + dy ** 2)
|
|
|
|
map_x = x_v + mask[0, :, :, 0] * dx / distance * (-lens_aperture * 100)
|
|
map_y = y_v + mask[0, :, :, 0] * dy / distance * (-lens_aperture * 100)
|
|
|
|
map_x = map_x.to(torch.float32).numpy()
|
|
map_y = map_y.to(torch.float32).numpy()
|
|
|
|
shifted_images = cv2_layer(images, lambda x: cv2.remap(x, map_x, map_y, cv2.INTER_LINEAR))
|
|
shifted_mask = cv2_layer(mask, lambda x: cv2.remap(x, map_x, map_y, cv2.INTER_LINEAR))
|
|
edited_images = cv2_layer(shifted_images, lambda x: cv2.stackBlur(x, (blur_intensity, blur_intensity)))
|
|
|
|
mask = torch.clamp(mask, 0.0, 1.0)
|
|
shifted_mask = torch.clamp(shifted_mask, 0.0, 1.0)
|
|
|
|
result = shifted_images * (1 - mask) + edited_images * mask
|
|
|
|
return result, shifted_mask
|
|
|
|
class LensZoomBurst:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"scale": ("FLOAT", {
|
|
"default": 1.5,
|
|
"min": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"samples": ("INT", {
|
|
"default": 100,
|
|
"min": 1,
|
|
}),
|
|
"position_x": ("FLOAT", {
|
|
"default": 0.5,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"position_y": ("FLOAT", {
|
|
"default": 0.5,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"rotation": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 360.0,
|
|
}),
|
|
"method": (["circle", "point"],),
|
|
"stabilization": (["true", "false"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "postprocessing/lens"
|
|
|
|
# noinspection PyUnresolvedReferences
|
|
def zoom_burst(
|
|
self,
|
|
image,
|
|
scale,
|
|
samples,
|
|
position,
|
|
rotation,
|
|
method,
|
|
stabilization,
|
|
):
|
|
if scale < 1.0:
|
|
raise ValueError("Parameter scale can't be smaller then initial image size.")
|
|
|
|
h, w = image.shape[:2]
|
|
|
|
x = np.arange(w)
|
|
y = np.arange(h)
|
|
|
|
xx, yy = np.meshgrid(x, y)
|
|
|
|
cx = int(w * position[0])
|
|
cy = int(h * position[1])
|
|
|
|
if method == "circle":
|
|
d = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)
|
|
max_d = np.sqrt((w / 2) ** 2 + (h / 2) ** 2)
|
|
|
|
map_x_up = (xx - d * (scale - 1.0) / max_d * (xx - cx) / samples).astype(np.float32)
|
|
map_y_up = (yy - d * (scale - 1.0) / max_d * (yy - cy) / samples).astype(np.float32)
|
|
|
|
map_x_down = (xx + d * (scale - 1.0) / max_d * (xx - cx) / samples).astype(np.float32)
|
|
map_y_down = (yy + d * (scale - 1.0) / max_d * (yy - cy) / samples).astype(np.float32)
|
|
elif method == "point":
|
|
map_x_up = (xx - (xx - cx) * (scale - 1.0) / samples).astype(np.float32)
|
|
map_y_up = (yy - (yy - cy) * (scale - 1.0) / samples).astype(np.float32)
|
|
|
|
map_x_down = (xx + (xx - cx) * (scale - 1.0) / samples).astype(np.float32)
|
|
map_y_down = (yy + (yy - cy) * (scale - 1.0) / samples).astype(np.float32)
|
|
else:
|
|
raise ValueError("Unsupported method.")
|
|
|
|
if rotation > 0.0:
|
|
angle_step = rotation / samples
|
|
|
|
rm_up = cv2.getRotationMatrix2D((cx, cy), angle_step, 1)
|
|
rm_down = cv2.getRotationMatrix2D((cx, cy), -angle_step, 1)
|
|
else:
|
|
vibration_angle = 1.0
|
|
vibration_step = vibration_angle / samples
|
|
|
|
rm_up_even = cv2.getRotationMatrix2D((cx, cy), vibration_step, 1)
|
|
rm_down_even = cv2.getRotationMatrix2D((cx, cy), -vibration_step, 1)
|
|
|
|
rm_up_odd = cv2.getRotationMatrix2D((cx, cy), -vibration_step, 1)
|
|
rm_down_odd = cv2.getRotationMatrix2D((cx, cy), vibration_step, 1)
|
|
|
|
for i in range(samples):
|
|
if stabilization:
|
|
tmp_up = cv2.remap(image, map_x_up, map_y_up, cv2.INTER_LINEAR)
|
|
tmp_down = cv2.remap(image, map_x_down, map_y_down, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
|
|
|
|
if rotation > 0.0:
|
|
tmp_up = cv2.warpAffine(tmp_up, rm_up, (w, h), borderMode=cv2.BORDER_REFLECT)
|
|
tmp_down = cv2.warpAffine(tmp_down, rm_down, (w, h), borderMode=cv2.BORDER_REFLECT)
|
|
else:
|
|
if i % 2 == 0:
|
|
tmp_up = cv2.warpAffine(tmp_up, rm_up_even, (w, h), borderMode=cv2.BORDER_REFLECT)
|
|
tmp_down = cv2.warpAffine(tmp_down, rm_down_even, (w, h), borderMode=cv2.BORDER_REFLECT)
|
|
else:
|
|
tmp_up = cv2.warpAffine(tmp_up, rm_up_odd, (w, h), borderMode=cv2.BORDER_REFLECT)
|
|
tmp_down = cv2.warpAffine(tmp_down, rm_down_odd, (w, h), borderMode=cv2.BORDER_REFLECT)
|
|
|
|
image = cv2.addWeighted(tmp_up, 0.5, tmp_down, 0.5, 0)
|
|
else:
|
|
tmp = cv2.remap(image, map_x_up, map_y_up, cv2.INTER_LINEAR)
|
|
|
|
if rotation > 0.0:
|
|
rm = cv2.getRotationMatrix2D((cx, cy), angle_step, 1)
|
|
tmp = cv2.warpAffine(tmp, rm, (w, h), borderMode=cv2.BORDER_REFLECT)
|
|
else:
|
|
if i % 2 == 0:
|
|
tmp = cv2.warpAffine(tmp, rm_up_even, (w, h), borderMode=cv2.BORDER_REFLECT)
|
|
else:
|
|
tmp = cv2.warpAffine(tmp, rm_up_odd, (w, h), borderMode=cv2.BORDER_REFLECT)
|
|
|
|
image = cv2.addWeighted(tmp, 0.5, image, 0.5, 0)
|
|
|
|
return image
|
|
|
|
def node(self, images, scale, samples, position_x, position_y, rotation, method, stabilization):
|
|
tensor = images.clone().detach()
|
|
|
|
return (cv2_layer(tensor, lambda x: self.zoom_burst(
|
|
x, scale, samples, (position_x, position_y), rotation, method, True if stabilization == "true" else False
|
|
)),)
|
|
|
|
class LUT:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image": ("IMAGE",),
|
|
"lut_image": ("IMAGE",),
|
|
"gamma_correction": (['True','False'],)}}
|
|
|
|
RETURN_TYPES = ('IMAGE',)
|
|
FUNCTION = 'apply_lut'
|
|
CATEGORY = 'Processing/color adjust'
|
|
OUTPUT_NODE = True
|
|
|
|
|
|
|
|
@apply_to_batch
|
|
def apply_lut(self, image, lut_image, gamma_correction):
|
|
img = tensor2pil(image)
|
|
lut_image = tensor2pil(lut_image)
|
|
|
|
if gamma_correction == 'True':
|
|
corrected_img = gamma_correction_pil(img, 1.0/2.2)
|
|
else:
|
|
corrected_img = img
|
|
filtered_image = apply_hald_clut(lut_image, corrected_img).convert("RGB")
|
|
#return (pil2tensor(filtered_image), )
|
|
return pil2tensor(filtered_image)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(self, hald_clut):
|
|
return (np.nan,)
|
|
|
|
from typing import Literal, Any
|
|
class Vignette:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"lens_shape": (["circle", "rectangle"],),
|
|
"lens_edge": (["around", "symmetric"],),
|
|
"lens_curvy": ("FLOAT", {
|
|
"default": 3.0,
|
|
"max": 15.0,
|
|
"step": 0.1,
|
|
}),
|
|
"lens_zoom": ("FLOAT", {
|
|
"default": 0.0,
|
|
"step": 0.1,
|
|
}),
|
|
"brightness": ("FLOAT", {
|
|
"default": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"saturation": ("FLOAT", {
|
|
"default": 0.5,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "node"
|
|
CATEGORY = "postprocessing/Effects"
|
|
|
|
def node(self, images, lens_shape, lens_edge, lens_curvy, lens_zoom, brightness, saturation):
|
|
tensor = images.clone().detach()
|
|
|
|
lens_zoom += 1
|
|
|
|
height, width = tensor[0, :, :, 0].shape
|
|
|
|
if lens_edge == "around":
|
|
mask = radialspace_2D((height, width), lens_curvy, lens_zoom, lens_shape).unsqueeze(0).unsqueeze(3)
|
|
elif lens_edge == "symmetric":
|
|
if height != width:
|
|
new_height = new_width = max(height, width)
|
|
crop_top_bottom = (new_height - height) // 2
|
|
crop_left_right = (new_width - width) // 2
|
|
|
|
mask = radialspace_2D((new_height, new_width), lens_curvy, lens_zoom, lens_shape)[
|
|
crop_top_bottom:-crop_top_bottom or None,
|
|
crop_left_right:-crop_left_right or None
|
|
].unsqueeze(0).unsqueeze(3)
|
|
else:
|
|
mask = radialspace_2D((height, width), lens_curvy, lens_zoom, lens_shape).unsqueeze(0).unsqueeze(3)
|
|
else:
|
|
raise ValueError("Not existing lens_edge parameter.")
|
|
|
|
tensor = tensor.permute(0, 3, 1, 2)
|
|
tensor[:, 0:3, :, :] = D.adjust_brightness(tensor[:, 0:3, :, :], brightness)
|
|
tensor[:, 0:3, :, :] = D.adjust_saturation(tensor[:, 0:3, :, :], saturation)
|
|
tensor = tensor.permute(0, 2, 3, 1)
|
|
|
|
result = images * (1 - mask) + tensor * mask
|
|
|
|
mask = mask.squeeze()
|
|
|
|
return result, mask
|
|
|
|
def gaussian_kernel(kernel_size: int, sigma: float):
|
|
x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij")
|
|
d = torch.sqrt(x * x + y * y)
|
|
g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
|
|
return g / g.sum()
|
|
|
|
def sort_span(span, sort_by, reverse_sorting):
|
|
if sort_by == 'H':
|
|
key = lambda x: x[1][0]
|
|
elif sort_by == 'S':
|
|
key = lambda x: x[1][1]
|
|
else:
|
|
key = lambda x: x[1][2]
|
|
|
|
span = sorted(span, key=key, reverse=reverse_sorting)
|
|
return [x[0] for x in span]
|
|
|
|
|
|
def find_spans(mask, span_limit=None):
|
|
spans = []
|
|
start = None
|
|
for i, value in enumerate(mask):
|
|
if value == 0 and start is None:
|
|
start = i
|
|
if value == 1 and start is not None:
|
|
span_length = i - start
|
|
if span_limit is None or span_length <= span_limit:
|
|
spans.append((start, i))
|
|
start = None
|
|
if start is not None:
|
|
span_length = len(mask) - start
|
|
if span_limit is None or span_length <= span_limit:
|
|
spans.append((start, len(mask)))
|
|
|
|
return spans
|
|
|
|
|
|
def pixel_sort(img, mask, horizontal_sort=False, span_limit=None, sort_by='H', reverse_sorting=False):
|
|
height, width, _ = img.shape
|
|
hsv_image = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32)
|
|
hsv_image[..., 0] /= 2.0 # Scale H channel to [0, 1] range
|
|
|
|
mask = np.where(mask > 0, 1, 0).astype(np.uint8)
|
|
|
|
# loop over the rows and replace contiguous bands of 1s
|
|
for i in range(height if horizontal_sort else width):
|
|
in_band = False
|
|
start = None
|
|
end = None
|
|
for j in range(width if horizontal_sort else height):
|
|
if (mask[i, j] if horizontal_sort else mask[j, i]) == 1:
|
|
if not in_band:
|
|
in_band = True
|
|
start = j
|
|
end = j
|
|
else:
|
|
if in_band:
|
|
for k in range(start+1, end):
|
|
if horizontal_sort:
|
|
mask[i, k] = 0
|
|
else:
|
|
mask[k, i] = 0
|
|
in_band = False
|
|
|
|
if in_band:
|
|
for k in range(start+1, end):
|
|
if horizontal_sort:
|
|
mask[i, k] = 0
|
|
else:
|
|
mask[k, i] = 0
|
|
|
|
sorted_image = np.zeros_like(img)
|
|
if horizontal_sort:
|
|
for y in range(height):
|
|
row_mask = mask[y]
|
|
spans = find_spans(row_mask, span_limit)
|
|
sorted_row = np.copy(img[y])
|
|
for start, end in spans:
|
|
span = [(img[y, x], hsv_image[y, x]) for x in range(start, end)]
|
|
sorted_span = sort_span(span, sort_by, reverse_sorting)
|
|
for i, pixel in enumerate(sorted_span):
|
|
sorted_row[start + i] = pixel
|
|
sorted_image[y] = sorted_row
|
|
else:
|
|
for x in range(width):
|
|
column_mask = mask[:, x]
|
|
spans = find_spans(column_mask, span_limit)
|
|
sorted_column = np.copy(img[:, x])
|
|
for start, end in spans:
|
|
span = [(img[y, x], hsv_image[y, x]) for y in range(start, end)]
|
|
sorted_span = sort_span(span, sort_by, reverse_sorting)
|
|
for i, pixel in enumerate(sorted_span):
|
|
sorted_column[start + i] = pixel
|
|
sorted_image[:, x] = sorted_column
|
|
|
|
return sorted_image
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"MX_Blend": Blend,
|
|
"MX_AlphaBlend": AlphaBlend,
|
|
"MX_Blur": Blur,
|
|
"MX_Canny": CannyEdgeMask,
|
|
"MX_ColorMatch": ColorMatch,
|
|
"MX_ChromaticAberration": ChromaticAberration,
|
|
"MX_ColorCorrect": ColorCorrect,
|
|
"MX_ColorTint": ColorTint,
|
|
"MX_Noise": FilmGrain,
|
|
"MX_Glow": Glow,
|
|
"MX_HSVThresholdMask": HSVThresholdMask,
|
|
"MX_KuwaharaBlur(Cartoon)": KuwaharaBlur,
|
|
"MX_Mosaic": Pixelize,
|
|
"MX_Posterize": Quantize,
|
|
"MX_SineWave": SineWave,
|
|
"MX_Solarize": Solarize,
|
|
"MX_LensBokeh": LensBokeh,
|
|
"MX_LensOpticAxis": LensOpticAxis,
|
|
"MX_LensZoomBurst": LensZoomBurst,
|
|
"MX_LUT":LUT,
|
|
"MX_Vignette": Vignette,
|
|
}
|