|
|
import cv2 |
|
|
import numpy as np |
|
|
import torch |
|
|
import torchvision.transforms.functional as F |
|
|
|
|
|
from .Utils import radialspace_1D, radialspace_2D, cv2_layer |
|
|
|
|
|
|
|
|
class ImageEffectsAdjustment: |
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(cls): |
|
|
return { |
|
|
"required": { |
|
|
"images": ("IMAGE",), |
|
|
"brightness": ("FLOAT", { |
|
|
"default": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
"contrast": ("FLOAT", { |
|
|
"default": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
"saturation": ("FLOAT", { |
|
|
"default": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
"hue": ("FLOAT", { |
|
|
"default": 0.5, |
|
|
"max": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
"gamma": ("FLOAT", { |
|
|
"default": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
"sharpness": ("FLOAT", { |
|
|
"default": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
"red": ("FLOAT", { |
|
|
"default": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
"green": ("FLOAT", { |
|
|
"default": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
"blue": ("FLOAT", { |
|
|
"default": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
}, |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("IMAGE",) |
|
|
FUNCTION = "node" |
|
|
CATEGORY = "image/effects" |
|
|
|
|
|
def node(self, images, brightness, contrast, saturation, hue, gamma, sharpness, red, green, blue): |
|
|
|
|
|
def apply(img): |
|
|
rgba = False |
|
|
|
|
|
if len(img[0, 0, :]) == 4: |
|
|
a = img[:, :, 3].unsqueeze(2) |
|
|
img = img[:, :, 0:3] |
|
|
rgba = True |
|
|
|
|
|
img = img.permute(2, 0, 1) |
|
|
|
|
|
if brightness != 1.0: |
|
|
img = F.adjust_brightness(img, brightness) |
|
|
|
|
|
if contrast != 1.0: |
|
|
img = F.adjust_contrast(img, contrast) |
|
|
|
|
|
if saturation != 1.0: |
|
|
img = F.adjust_saturation(img, saturation) |
|
|
|
|
|
if hue != 0.5: |
|
|
img = F.adjust_hue(img, hue - 0.5) |
|
|
|
|
|
if gamma != 1.0: |
|
|
img = F.adjust_gamma(img, gamma) |
|
|
|
|
|
if sharpness != 1.0: |
|
|
img = F.adjust_sharpness(img, sharpness) |
|
|
|
|
|
img = img.permute(1, 2, 0) |
|
|
|
|
|
r, g, b = torch.chunk(img, 3, dim=2) |
|
|
|
|
|
r = torch.clamp(r * red, 0, 1) |
|
|
g = torch.clamp(g * green, 0, 1) |
|
|
b = torch.clamp(b * blue, 0, 1) |
|
|
|
|
|
if rgba: |
|
|
return torch.cat([r, g, b, a], dim=2) |
|
|
else: |
|
|
return torch.cat([r, g, b], dim=2) |
|
|
|
|
|
return (torch.stack([ |
|
|
apply(images[i]) for i in range(len(images)) |
|
|
]),) |
|
|
|
|
|
|
|
|
class ImageEffectsGrayscale: |
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(cls): |
|
|
return { |
|
|
"required": { |
|
|
"images": ("IMAGE",), |
|
|
}, |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("IMAGE",) |
|
|
FUNCTION = "node" |
|
|
CATEGORY = "image/effects" |
|
|
|
|
|
def node(self, images): |
|
|
def apply(image): |
|
|
tensor = image.clone().detach() |
|
|
grayscale_tensor = torch.mean(tensor, dim=2, keepdim=True) |
|
|
|
|
|
return torch.cat([grayscale_tensor] * 3, dim=2) |
|
|
|
|
|
return (torch.stack([ |
|
|
apply(images[i]) for i in range(len(images)) |
|
|
]),) |
|
|
|
|
|
|
|
|
class ImageEffectsNegative: |
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(cls): |
|
|
return { |
|
|
"required": { |
|
|
"images": ("IMAGE",), |
|
|
}, |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("IMAGE",) |
|
|
FUNCTION = "node" |
|
|
CATEGORY = "image/effects" |
|
|
|
|
|
def node(self, images): |
|
|
tensor = images.clone().detach() |
|
|
tensor[:, :, :, 0:3] = 1.0 - tensor[:, :, :, 0:3] |
|
|
|
|
|
return (tensor,) |
|
|
|
|
|
|
|
|
class ImageEffectsSepia: |
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(cls): |
|
|
return { |
|
|
"required": { |
|
|
"images": ("IMAGE",), |
|
|
}, |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("IMAGE",) |
|
|
FUNCTION = "node" |
|
|
CATEGORY = "image/effects" |
|
|
|
|
|
def node(self, images): |
|
|
tensor = images.clone().detach() |
|
|
|
|
|
sepia_mask = torch.tensor([[0.393, 0.349, 0.272], |
|
|
[0.769, 0.686, 0.534], |
|
|
[0.189, 0.168, 0.131]]) |
|
|
|
|
|
tensor[:, :, :, 0:3] = torch.stack([ |
|
|
torch.matmul(tensor[i, :, :, 0:3], sepia_mask) for i in range(len(tensor)) |
|
|
]) |
|
|
|
|
|
return (tensor,) |
|
|
|
|
|
|
|
|
class ImageEffectsLensZoomBurst: |
|
|
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 = "image/effects/lens" |
|
|
|
|
|
|
|
|
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 ImageEffectsLensChromaticAberration: |
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
@classmethod |
|
|
def INPUT_TYPES(cls): |
|
|
return { |
|
|
"required": { |
|
|
"images": ("IMAGE",), |
|
|
"shift": ("INT", { |
|
|
"default": 10, |
|
|
"step": 1, |
|
|
}), |
|
|
"method": (["reflect", "edge", "constant"],), |
|
|
"shift_type": ("INT", { |
|
|
"default": 1, |
|
|
"min": 1, |
|
|
"max": 4, |
|
|
"step": 1, |
|
|
}), |
|
|
"mixing_type": ("INT", { |
|
|
"default": 1, |
|
|
"min": 1, |
|
|
"max": 4, |
|
|
"step": 1, |
|
|
}), |
|
|
"transpose": (["none", "rotate", "reflect"],), |
|
|
"colors": (["rb", "rg", "gb"],), |
|
|
"lens_curvy": ("FLOAT", { |
|
|
"default": 1.0, |
|
|
"max": 15.0, |
|
|
"step": 0.1, |
|
|
}), |
|
|
}, |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("IMAGE",) |
|
|
FUNCTION = "node" |
|
|
CATEGORY = "image/effects/lens" |
|
|
|
|
|
def node(self, images, shift, method, shift_type, mixing_type, transpose, colors, lens_curvy): |
|
|
|
|
|
def apply(image): |
|
|
img = image.clone().detach() |
|
|
|
|
|
if transpose == "rotate": |
|
|
img = img.permute(1, 0, 2) |
|
|
elif transpose == "none" or transpose == "reflect": |
|
|
pass |
|
|
else: |
|
|
raise ValueError("Not existing reverse.") |
|
|
|
|
|
r, g, b = img[:, :, 0:3].split(1, 2) |
|
|
height, width = img[:, :, 0].shape |
|
|
|
|
|
def get_space(size, min_val, max_val): |
|
|
size += shift * 2 |
|
|
|
|
|
if min_val == max_val: |
|
|
return torch.linspace(min_val, max_val, size) |
|
|
else: |
|
|
return radialspace_1D(size, lens_curvy, 1.0, min_val, max_val) |
|
|
|
|
|
if shift_type == 1: |
|
|
f_shifts = get_space(height, -shift, shift) |
|
|
|
|
|
if transpose == "reflect": |
|
|
t_shifts = get_space(width, -shift, shift) |
|
|
elif shift_type == 2: |
|
|
f_shifts = get_space(height, 0, shift) |
|
|
f_shifts = torch.flip(f_shifts, dims=(0,)) |
|
|
|
|
|
if transpose == "reflect": |
|
|
t_shifts = get_space(width, 0, shift) |
|
|
t_shifts = torch.flip(t_shifts, dims=(0,)) |
|
|
elif shift_type == 3: |
|
|
f_shifts = get_space(height, 0, shift) |
|
|
|
|
|
if transpose == "reflect": |
|
|
t_shifts = get_space(width, 0, shift) |
|
|
elif shift_type == 4: |
|
|
f_shifts = get_space(height, shift, shift) |
|
|
|
|
|
if transpose == "reflect": |
|
|
t_shifts = get_space(width, shift, shift) |
|
|
else: |
|
|
raise ValueError("Not existing shift_type.") |
|
|
|
|
|
if mixing_type == 1: |
|
|
f_shifts = f_shifts |
|
|
s_shifts = -f_shifts |
|
|
|
|
|
if transpose == "reflect": |
|
|
t_shifts = t_shifts |
|
|
c_shifts = -t_shifts |
|
|
elif mixing_type == 2: |
|
|
f_shifts = -f_shifts |
|
|
s_shifts = f_shifts |
|
|
|
|
|
if transpose == "reflect": |
|
|
t_shifts = -t_shifts |
|
|
c_shifts = t_shifts |
|
|
elif mixing_type == 3: |
|
|
f_shifts = f_shifts |
|
|
s_shifts = f_shifts |
|
|
|
|
|
if transpose == "reflect": |
|
|
t_shifts = t_shifts |
|
|
c_shifts = t_shifts |
|
|
elif mixing_type == 4: |
|
|
f_shifts = -f_shifts |
|
|
s_shifts = -f_shifts |
|
|
|
|
|
if transpose == "reflect": |
|
|
t_shifts = -t_shifts |
|
|
c_shifts = -t_shifts |
|
|
else: |
|
|
raise ValueError("Not existing mixing_type.") |
|
|
|
|
|
if colors == "rb": |
|
|
def cat(f_value, s_value): |
|
|
return torch.cat([f_value, g, s_value], 2) |
|
|
|
|
|
f = r |
|
|
s = b |
|
|
elif colors == "rg": |
|
|
def cat(f_value, s_value): |
|
|
return torch.cat([f_value, s_value, b], 2) |
|
|
|
|
|
f = r |
|
|
s = g |
|
|
elif colors == "gb": |
|
|
def cat(f_value, s_value): |
|
|
return torch.cat([r, f_value, s_value], 2) |
|
|
|
|
|
f = g |
|
|
s = b |
|
|
else: |
|
|
raise ValueError("Not existing colors.") |
|
|
|
|
|
f_pad = F.pad(f.squeeze(), [shift, shift], padding_mode=method).unsqueeze(2) |
|
|
s_pad = F.pad(s.squeeze(), [shift, shift], padding_mode=method).unsqueeze(2) |
|
|
|
|
|
f_shifted = torch.zeros_like(f_pad) |
|
|
s_shifted = torch.zeros_like(s_pad) |
|
|
|
|
|
for i in range(height + (shift * 2)): |
|
|
f_shifted[i] = torch.roll(f_pad[i], shifts=int(f_shifts[i]), dims=0) |
|
|
s_shifted[i] = torch.roll(s_pad[i], shifts=int(s_shifts[i]), dims=0) |
|
|
|
|
|
if transpose == "reflect": |
|
|
for i in range(width + (shift * 2)): |
|
|
f_shifted[:, i] = torch.roll(f_shifted[:, i], shifts=int(t_shifts[i]), dims=0) |
|
|
s_shifted[:, i] = torch.roll(s_shifted[:, i], shifts=int(c_shifts[i]), dims=0) |
|
|
|
|
|
f_result = f_shifted[shift:-shift, shift:-shift, :] |
|
|
s_result = s_shifted[shift:-shift, shift:-shift, :] |
|
|
|
|
|
img[:, :, 0:3] = cat(f_result, s_result) |
|
|
|
|
|
if transpose == "rotate": |
|
|
img = img.permute(1, 0, 2) |
|
|
elif transpose == "none" or transpose == "reflect": |
|
|
pass |
|
|
else: |
|
|
raise ValueError("Not existing reverse.") |
|
|
|
|
|
return img |
|
|
|
|
|
return (torch.stack([ |
|
|
apply(images[i]) for i in range(len(images)) |
|
|
]),) |
|
|
|
|
|
|
|
|
class ImageEffectsLensBokeh: |
|
|
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 = "image/effects/lens" |
|
|
|
|
|
|
|
|
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 ImageEffectsLensOpticAxis: |
|
|
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 = "image/effects/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 ImageEffectsLensVignette: |
|
|
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.25, |
|
|
"max": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
"saturation": ("FLOAT", { |
|
|
"default": 0.5, |
|
|
"max": 1.0, |
|
|
"step": 0.01 |
|
|
}), |
|
|
}, |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK") |
|
|
FUNCTION = "node" |
|
|
CATEGORY = "image/effects/lens" |
|
|
|
|
|
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, :, :] = F.adjust_brightness(tensor[:, 0:3, :, :], brightness) |
|
|
tensor[:, 0:3, :, :] = F.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 |
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = { |
|
|
"ImageEffectsAdjustment": ImageEffectsAdjustment, |
|
|
"ImageEffectsGrayscale": ImageEffectsGrayscale, |
|
|
"ImageEffectsNegative": ImageEffectsNegative, |
|
|
"ImageEffectsSepia": ImageEffectsSepia, |
|
|
"ImageEffectsLensZoomBurst": ImageEffectsLensZoomBurst, |
|
|
"ImageEffectsLensChromaticAberration": ImageEffectsLensChromaticAberration, |
|
|
"ImageEffectsLensBokeh": ImageEffectsLensBokeh, |
|
|
"ImageEffectsLensOpticAxis": ImageEffectsLensOpticAxis, |
|
|
"ImageEffectsLensVignette": ImageEffectsLensVignette |
|
|
} |
|
|
|
