nodes/EllangoK-ComfyUI-post-processing-nodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add Sepia Effect

Browse Source
This commit is contained in:
EllangoK
2023-05-04 12:04:22 -04:00
parent 726da7523d
commit 33a525b059
4 changed files with 90 additions and 363 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -34,11 +34,11 @@ Both images have the workflow attached, and are included with the repo. Feel fre
- $\color{#00A7B5}\textbf{PixelSort:}$ Rearranges the pixels in the input image based on their values, and input mask. Creates a cool glitch like effect.
- Pixelize: Applies a pixelization effect, simulating the reducing of resolution
- $\color{#00A7B5}\textbf{Quantize:}$ Set and dither the amount of colors in an image from 0-256, reducing color information
- Sepia: Applies a mellow tone mapping, yielding an archival or vintage appearance
- Sharpen: Enhances the details in an image by applying a sharpening filter
- $\color{#00A7B5}\textbf{Solarize:}$ Inverts image colors based on a threshold for a striking, high-contrast effect
- Vignette: Applies a vignette effect, putting the corners of the image in shadow
$\color{#00A7B5}\textbf{Bolded Color Nodes}$ are my personal favorites, and highly recommended to expirement with
</details>
combine_files.py
+1 -5
View File
@@ -1,9 +1,5 @@
from collections import OrderedDict
from pathlib import Path
import sys
import os
import glob
import ast
import argparse
ignore_dirs = ["old"]
@@ -12,7 +8,7 @@ def get_python_files(path, recursive=False, args=None):
search_pattern = "**/*.py" if recursive else "*.py"
def should_include(file):
if file.is_file() and not file.name.startswith("combine") and not args.output in str(file):
if file.is_file() and not file.name.startswith("combine") and not args.output in str(file) and not file.name.startswith("__init__"):
for ignore_dir in ignore_dirs:
if ignore_dir in str(file.parent):
return False
post_processing/sepia.py
+41
View File
@@ -0,0 +1,41 @@
import torch
class Sepia:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"strength": ("FLOAT", {
"default": 1.0,
"min": 0.1,
"max": 1.0,
"step": 0.1
}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "sepia"
CATEGORY = "postprocessing"
def sepia(self, image: torch.Tensor, strength: float):
if strength == 0:
return (image,)
sepia_weights = torch.tensor([0.2989, 0.5870, 0.1140]).view(1, 1, 1, 3).to(image.device)
sepia_filter = torch.tensor([1.0, 0.8, 0.6]).view(1, 1, 1, 3).to(image.device)
grayscale = torch.sum(image * sepia_weights, dim=-1, keepdim=True)
sepia = grayscale * sepia_filter
result = sepia * strength + image * (1 - strength)
return (result,)
NODE_CLASS_MAPPINGS = {
"Sepia": Sepia
}
post_processing_nodes.py
+46 -356
View File
@@ -5,8 +5,6 @@ import numpy as np
from PIL import Image, ImageEnhance
import multiprocessing as mp
from PIL import Image
import time
import random
class ArithmeticBlend:
@@ -673,6 +671,7 @@ class KuwaharaBlur:
"max": 31,
"step": 1
}),
"method": (["mean", "gaussian"],),
},
}
@@ -681,7 +680,7 @@ class KuwaharaBlur:
CATEGORY = "postprocessing"
def apply_kuwahara_filter(self, image: np.ndarray, blur_radius: int):
def apply_kuwahara_filter(self, image: np.ndarray, blur_radius: int, method: str):
if blur_radius == 0:
return (image,)
@@ -692,21 +691,18 @@ class KuwaharaBlur:
image = image[b].cpu().numpy() * 255.0
image = image.astype(np.uint8)
out[b] = torch.from_numpy(kuwahara(image, method="gaussian", radius=blur_radius)) / 255.0
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, grayconv=cv2.COLOR_BGR2GRAY, image_2d=None):
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)
image_2d = image_2d.astype(image.dtype, copy=False) if image_2d is not None else None
avgs = np.empty((4, *image.shape), dtype=image.dtype)
stddevs = np.empty((4, *image.shape[:2]), dtype=image.dtype)
if image_2d is None:
image_2d = cv2.cvtColor(orig_img, grayconv).astype(image.dtype, copy=False)
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
@@ -722,7 +718,10 @@ def kuwahara(orig_img, method="mean", radius=3, sigma=None, grayconv=cv2.COLOR_B
shift = [(0, 0), (0, radius), (radius, 0), (radius, radius)]
for k in range(4):
kx, ky = kxy, kxy if method == "mean" else klr[kindexes[k]]
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])
@@ -985,6 +984,42 @@ class Quantize:
return (result,)
class Sepia:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"strength": ("FLOAT", {
"default": 1.0,
"min": 0.1,
"max": 1.0,
"step": 0.1
}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "sepia"
CATEGORY = "postprocessing"
def sepia(self, image: torch.Tensor, strength: float):
if strength == 0:
return (image,)
sepia_weights = torch.tensor([0.2989, 0.5870, 0.1140]).view(1, 1, 1, 3).to(image.device)
sepia_filter = torch.tensor([1.0, 0.8, 0.6]).view(1, 1, 1, 3).to(image.device)
grayscale = torch.sum(image * sepia_weights, dim=-1, keepdim=True)
sepia = grayscale * sepia_filter
result = sepia * strength + image * (1 - strength)
return (result,)
class Sharpen:
def __init__(self):
pass
@@ -1104,265 +1139,6 @@ class Vignette:
return (vignette_image,)
class ElectroShock:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"glow_intensity": ("INT", {"default": 50, "min": 0, "max": 100, "step": 1}),
"line_frequency": ("INT", {"default": 25, "min": 0, "max": 100, "step": 1}),
"line_thickness": ("INT", {"default": 2, "min": 1, "max": 10, "step": 1}),
"random_seed": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "electro_shock"
CATEGORY = "effects"
def midpoint_displacement(self, x1, y1, x2, y2, displacement, mask, line_thickness):
if abs(x2 - x1) < 2 and abs(y2 - y1) < 2:
return
mid_x = (x1 + x2) // 2
mid_y = (y1 + y2) // 2
mid_x += int(random.uniform(-displacement, displacement))
mid_y += int(random.uniform(-displacement, displacement))
cv2.line(mask, (x1, y1), (mid_x, mid_y), 255, line_thickness)
cv2.line(mask, (mid_x, mid_y), (x2, y2), 255, line_thickness)
self.midpoint_displacement(x1, y1, mid_x, mid_y, displacement / 2, mask, line_thickness)
self.midpoint_displacement(mid_x, mid_y, x2, y2, displacement / 2, mask, line_thickness)
def electro_shock(self, image: torch.Tensor, glow_intensity: int, line_frequency: int, line_thickness: int, random_seed: int = None):
if random_seed is not None:
random.seed(random_seed)
np.random.seed(random_seed)
line_color = [255, 255, 255]
batch_size, height, width, _ = image.shape
result = torch.zeros_like(image)
for b in range(batch_size):
tensor_image = image[b]
img = (tensor_image * 255).to(torch.uint8).numpy()
# Apply the ElectroShock effect using OpenCV functions
mask = np.zeros((height, width), np.uint8)
num_lines = int(line_frequency * (height * width) / 100000)
initial_displacement = int(height / 8)
for _ in range(num_lines):
x1, y1 = random.randint(0, width - 1), random.randint(0, height - 1)
x2, y2 = random.randint(0, width - 1), random.randint(0, height - 1)
self.midpoint_displacement(x1, y1, x2, y2, initial_displacement, mask, line_thickness)
# Apply glow effect
glow_radius = int(glow_intensity * 0.1)
mask_blurred = cv2.GaussianBlur(mask, (glow_radius * 2 + 1, glow_radius * 2 + 1), 0)
# Add glow to the original image
colored_mask = cv2.cvtColor(mask_blurred, cv2.COLOR_GRAY2BGR)
colored_mask[np.where((colored_mask == [255, 255, 255]).all(axis=2))] = line_color
electro_shock_img = cv2.addWeighted(img, 1, colored_mask, glow_intensity / 100, 0)
electro_shock_array = torch.tensor(electro_shock_img).float() / 255
result[b] = electro_shock_array
return (result,)
class KuwaharaFilter:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"blur_radius": ("INT", {
"default": 1,
"min": 0,
"max": 15,
"step": 1
}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "apply_kuwahara_filter"
CATEGORY = "postprocessing"
def apply_kuwahara_filter(self, image: np.ndarray, blur_radius: int):
if blur_radius == 0:
return (image,)
kernel_size = blur_radius * 2 + 1
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_filter_rgb(image, kernel_size)) / 255.0
return (out,)
def kuwahara_filter_rgb(img, kernel_size):
b, g, r = cv2.split(img)
b_filtered, g_filtered, r_filtered = apply_filter((b, kernel_size)), apply_filter((g, kernel_size)), apply_filter((r, kernel_size))
out = cv2.merge((b_filtered, g_filtered, r_filtered))
return out
def apply_filter(args):
channel, kernel_size = args
return kuwahara_filter(channel, kernel_size)
def kuwahara_filter(img, kernel_size):
# Pad the image to handle borders
pad_size = kernel_size // 2
img_padded = cv2.copyMakeBorder(img, pad_size, pad_size, pad_size, pad_size, cv2.BORDER_REFLECT)
# Initialize output image
h, w = img.shape[:2]
out = np.zeros_like(img)
# Apply Kuwahara filter to each pixel
for i in range(pad_size, h + pad_size):
for j in range(pad_size, w + pad_size):
# Divide the image into 4 overlapping square regions
regions = [
img_padded[i-pad_size:i+pad_size+1, j-pad_size:j+pad_size+1],
img_padded[i-pad_size:i+pad_size+1, j:j+kernel_size+1],
img_padded[i:i+kernel_size+1, j-pad_size:j+pad_size+1],
img_padded[i:i+kernel_size+1, j:j+kernel_size+1]
]
# Compute mean and variance of each region
means = [np.mean(region) for region in regions]
variances = [np.var(region) for region in regions]
# Choose the region with the smallest variance as the output value
min_var_index = np.argmin(variances)
out[i-pad_size, j-pad_size] = means[min_var_index]
return out
class Liquidify:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"viscosity": ("INT", {
"default": 10,
"min": 0,
"max": 20,
"step": 1
}),
"turbulence": ("FLOAT", {
"default": 1.0,
"min": 0.0,
"max": 2.0,
"step": 0.1
}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "liquidify"
CATEGORY = "postprocessing"
def liquidify(image: torch.Tensor, viscosity: int, turbulence: float):
image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
n, c, h, w = image.size()
grid_x, grid_y = torch.meshgrid(torch.arange(h), torch.arange(w))
grid_x = grid_x.to(image.device)
grid_y = grid_y.to(image.device)
displacement = torch.randn(n, 2, h, w).to(image.device)
displacement = F.gaussian_blur(displacement, kernel_size=viscosity, sigma=turbulence)
flow_x = torch.clamp(grid_x + displacement[:, 0], 0, w - 1).unsqueeze(1) - grid_x.unsqueeze(0)
flow_y = torch.clamp(grid_y + displacement[:, 1], 0, h - 1).unsqueeze(1) - grid_y.unsqueeze(0)
warped = F.grid_sample(image, torch.stack((flow_x, flow_y), dim=1), padding_mode='border')
warped = warped.permute(0, 2, 3, 1) # Back to (B, H, W, C)
return (warped,)
class StippleEffect:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"dot_size": ("FLOAT", {
"default": 1.0,
"min": 0.1,
"max": 5.0,
"step": 0.1
}),
"density": ("FLOAT", {
"default": 1.0,
"min": 0.1,
"max": 5.0,
"step": 0.1
}),
"intensity": ("FLOAT", {
"default": 1.0,
"min": 0.1,
"max": 5.0,
"step": 0.1
}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "stipple_effect"
CATEGORY = "postprocessing"
def stipple_effect(self, image: torch.Tensor, dot_size: float, density: float, intensity: float):
def create_dot_pattern(dot_size, intensity):
dot_pattern = torch.ones((1, 1, int(dot_size), int(dot_size))) * intensity
return dot_pattern
x = image.permute(0, 3, 1, 2)
gray_image = x.mean(dim=1, keepdim=True)
dot_pattern = create_dot_pattern(dot_size, intensity)
stippled_image = torch.nn.functional.conv2d(gray_image, dot_pattern, stride=int(dot_size), groups=1)
stippled_image = torch.clamp(stippled_image, 0, 1)
output = stippled_image.expand(-1, 3, -1, -1)
output = output.permute(0, 2, 3, 1)
return (output,)
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)
@@ -1460,89 +1236,6 @@ def pixel_sort(img, mask, horizontal_sort=False, span_limit=None, sort_by='H', r
return sorted_image
def kuwahara_filter(img, kernel_size):
# Pad the image to handle borders
pad_size = kernel_size // 2
img_padded = cv2.copyMakeBorder(img, pad_size, pad_size, pad_size, pad_size, cv2.BORDER_REFLECT)
# Initialize output image
h, w = img.shape[:2]
out = np.zeros_like(img)
# Apply Kuwahara filter to each pixel
for i in range(pad_size, h + pad_size):
for j in range(pad_size, w + pad_size):
# Divide the image into 4 overlapping square regions
regions = [
img_padded[i-pad_size:i+pad_size+1, j-pad_size:j+pad_size+1],
img_padded[i-pad_size:i+pad_size+1, j:j+kernel_size+1],
img_padded[i:i+kernel_size+1, j-pad_size:j+pad_size+1],
img_padded[i:i+kernel_size+1, j:j+kernel_size+1]
]
# Compute mean and variance of each region
means = [np.mean(region) for region in regions]
variances = [np.var(region) for region in regions]
# Choose the region with the smallest variance as the output value
min_var_index = np.argmin(variances)
out[i-pad_size, j-pad_size] = means[min_var_index]
return out
def kuwahara_filter_rgb(img, kernel_size):
# Split the image into color channels
b, g, r = cv2.split(img)
# Apply the filter to each channel
b_filtered = kuwahara_filter(b, kernel_size)
g_filtered = kuwahara_filter(g, kernel_size)
r_filtered = kuwahara_filter(r, kernel_size)
# Merge the filtered channels back into an RGB image
out = cv2.merge((b_filtered, g_filtered, r_filtered))
return out
def apply_filter(args):
channel, kernel_size = args
return kuwahara_filter(channel, kernel_size)
def kuwahara_filter_rgb_multiprocessing(img, kernel_size):
# Split the image into color channels
b, g, r = cv2.split(img)
# Function to apply the filter to a channel
# Create a multiprocessing Pool with 3 processes
with mp.Pool(3) as pool:
# Map the apply_filter function to the channels
b_filtered, g_filtered, r_filtered = pool.map(apply_filter, ((b, kernel_size), (g, kernel_size), (r, kernel_size)))
# Merge the filtered channels back into an RGB image
out = cv2.merge((b_filtered, g_filtered, r_filtered))
return out
if __name__ == '__main__':
img = cv2.imread('test.png')
start_time = time.time()
# Apply Kuwahara filter with kernel size of 5
out = kuwahara_filter_rgb_multiprocessing(img, kernel_size=5)
end_time = time.time()
print(f"Time elapsed: {end_time - start_time:.5f} seconds")
# Display output image
cv2.imshow('output_image', out)
cv2.waitKey(0)
cv2.destroyAllWindows()
NODE_CLASS_MAPPINGS = {
"ArithmeticBlend": ArithmeticBlend,
"Blend": Blend,
@@ -1560,11 +1253,8 @@ NODE_CLASS_MAPPINGS = {
"PixelSort": PixelSort,
"Pixelize": Pixelize,
"Quantize": Quantize,
"Sepia": Sepia,
"Sharpen": Sharpen,
"Solarize": Solarize,
"Vignette": Vignette,
"ElectroShock": ElectroShock,
"KuwaharaFilter": KuwaharaFilter,
"Liquidify": Liquidify,
"StippleEffect": StippleEffect,
}