From 33a525b0591d5cf0c7386714dc125f5cb1251763 Mon Sep 17 00:00:00 2001 From: EllangoK Date: Thu, 4 May 2023 12:04:22 -0400 Subject: [PATCH] Add Sepia Effect --- README.md | 4 +- combine_files.py | 6 +- post_processing/sepia.py | 41 ++++ post_processing_nodes.py | 402 +++++---------------------------------- 4 files changed, 90 insertions(+), 363 deletions(-) create mode 100644 post_processing/sepia.py diff --git a/README.md b/README.md index b8b3755..ea361d6 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/combine_files.py b/combine_files.py index a7b87f9..bf5b58b 100644 --- a/combine_files.py +++ b/combine_files.py @@ -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 diff --git a/post_processing/sepia.py b/post_processing/sepia.py new file mode 100644 index 0000000..521be94 --- /dev/null +++ b/post_processing/sepia.py @@ -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 +} diff --git a/post_processing_nodes.py b/post_processing_nodes.py index e581eb8..1013a4c 100644 --- a/post_processing_nodes.py +++ b/post_processing_nodes.py @@ -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, }