From 038a36d25e2be245ed63edb642c674781cf69240 Mon Sep 17 00:00:00 2001 From: EllangoK Date: Mon, 3 Apr 2023 17:01:00 -0400 Subject: [PATCH] adds PencilSketch effect, mimics a drawn sketch also extracts out gaussian_kernel func --- README.md | 1 + post_processing/blur.py | 14 ++-- post_processing/glow.py | 16 ++--- post_processing/pencil_sketch.py | 97 +++++++++++++++++++++++++++ post_processing_nodes.py | 111 ++++++++++++++++++++++++++----- 5 files changed, 209 insertions(+), 30 deletions(-) create mode 100644 post_processing/pencil_sketch.py diff --git a/README.md b/README.md index 7984d9d..da1a90a 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ A collection of post processing nodes for [ComfyUI](https://github.com/comfyanon - DodgeAndBurn: Adjusts image brightness using dodge and burn effects based on a mask and intensity. - FilmGrain: Adds a film grain effect to the image, along with options to control the temperature, and vignetting - Glow: Applies a blur with a specified radius and then blends it with the original image. Creates a nice glowing effect. + - PencilSketch: Converts an image into a hand-drawn pencil sketch style. - 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 - Quantize: Set and dither the amount of colors in an image from 0-256, reducing color information diff --git a/post_processing/blur.py b/post_processing/blur.py index 7553292..221db15 100644 --- a/post_processing/blur.py +++ b/post_processing/blur.py @@ -30,12 +30,6 @@ class Blur: CATEGORY = "postprocessing" - def gaussian_kernel(self, 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 blur(self, image: torch.Tensor, blur_radius: int, sigma: float): if blur_radius == 0: return (image,) @@ -43,7 +37,7 @@ class Blur: batch_size, height, width, channels = image.shape kernel_size = blur_radius * 2 + 1 - kernel = self.gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(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) @@ -54,3 +48,9 @@ class Blur: NODE_CLASS_MAPPINGS = { "Blur": Blur } + +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() diff --git a/post_processing/glow.py b/post_processing/glow.py index 40fb7eb..28bab4c 100644 --- a/post_processing/glow.py +++ b/post_processing/glow.py @@ -36,17 +36,11 @@ class Glow: glowing_image = torch.clamp(glowing_image, 0, 1) return (glowing_image,) - def gaussian_kernel(self, kernel_size: int): - sigma = (kernel_size - 1) / 6 - x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size)) - d = torch.sqrt(x * x + y * y) - g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) - return g / g.sum() - def gaussian_blur(self, image: torch.Tensor, kernel_size: int): batch_size, height, width, channels = image.shape - kernel = self.gaussian_kernel(kernel_size).repeat(channels, 1, 1).unsqueeze(1) + 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) @@ -60,3 +54,9 @@ class Glow: NODE_CLASS_MAPPINGS = { "Glow": Glow, } + +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() diff --git a/post_processing/pencil_sketch.py b/post_processing/pencil_sketch.py new file mode 100644 index 0000000..8640c9d --- /dev/null +++ b/post_processing/pencil_sketch.py @@ -0,0 +1,97 @@ +import torch +import torch.nn.functional as F + +class PencilSketch: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "blur_radius": ("INT", { + "default": 5, + "min": 1, + "max": 31, + "step": 1 + }), + "sharpen_alpha": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 10.0, + "step": 0.1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply_sketch" + + CATEGORY = "postprocessing" + + def apply_sketch(self, image: torch.Tensor, blur_radius: int = 5, sharpen_alpha: float = 1): + image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) + + grayscale = image.mean(dim=1, keepdim=True) + grayscale = grayscale.repeat(1, 3, 1, 1) + inverted = 1 - grayscale + + blur_sigma = blur_radius / 3 + blurred = self.gaussian_blur(inverted, blur_radius, blur_sigma) + + final_image = self.dodge(blurred, grayscale) + + if sharpen_alpha != 0.0: + final_image = self.sharpen(final_image, 1, sharpen_alpha) + + final_image = final_image.permute(0, 2, 3, 1) # Back to (B, H, W, C) + + return (final_image,) + + def dodge(self, front: torch.Tensor, back: torch.Tensor) -> torch.Tensor: + result = back / (1 - front + 1e-7) + result = torch.clamp(result, 0, 1) + return result + + def gaussian_blur(self, image: torch.Tensor, blur_radius: int, sigma: float): + if blur_radius == 0: + return image + + batch_size, channels, height, width = image.shape + + kernel_size = blur_radius * 2 + 1 + kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1) + + blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels) + + return blurred + + def sharpen(self, image: torch.Tensor, blur_radius: int, alpha: float): + if blur_radius == 0: + return image + + batch_size, channels, height, width = image.shape + + kernel_size = blur_radius * 2 + 1 + kernel = torch.ones((kernel_size, kernel_size), dtype=torch.float32) * -1 + center = kernel_size // 2 + kernel[center, center] = kernel_size**2 + kernel *= alpha + kernel = kernel.repeat(channels, 1, 1).unsqueeze(1) + + sharpened = F.conv2d(image, kernel, padding=center, groups=channels) + + result = torch.clamp(sharpened, 0, 1) + + return result + +NODE_CLASS_MAPPINGS = { + "PencilSketch": PencilSketch, +} + +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() diff --git a/post_processing_nodes.py b/post_processing_nodes.py index b5c6ac7..054ff84 100644 --- a/post_processing_nodes.py +++ b/post_processing_nodes.py @@ -3,6 +3,7 @@ import torch.nn.functional as F import cv2 import numpy as np from PIL import Image, ImageEnhance +from PIL import Image class ArithmeticBlend: @@ -149,12 +150,6 @@ class Blur: CATEGORY = "postprocessing" - def gaussian_kernel(self, 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 blur(self, image: torch.Tensor, blur_radius: int, sigma: float): if blur_radius == 0: return (image,) @@ -162,7 +157,7 @@ class Blur: batch_size, height, width, channels = image.shape kernel_size = blur_radius * 2 + 1 - kernel = self.gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(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) @@ -593,17 +588,11 @@ class Glow: glowing_image = torch.clamp(glowing_image, 0, 1) return (glowing_image,) - def gaussian_kernel(self, kernel_size: int): - sigma = (kernel_size - 1) / 6 - x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size)) - d = torch.sqrt(x * x + y * y) - g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) - return g / g.sum() - def gaussian_blur(self, image: torch.Tensor, kernel_size: int): batch_size, height, width, channels = image.shape - kernel = self.gaussian_kernel(kernel_size).repeat(channels, 1, 1).unsqueeze(1) + 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) @@ -614,6 +603,91 @@ class Glow: def add_glow(self, img, blurred_img, intensity): return img + blurred_img * intensity +class PencilSketch: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "blur_radius": ("INT", { + "default": 5, + "min": 1, + "max": 31, + "step": 1 + }), + "sharpen_alpha": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 10.0, + "step": 0.1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply_sketch" + + CATEGORY = "postprocessing" + + def apply_sketch(self, image: torch.Tensor, blur_radius: int = 5, sharpen_alpha: float = 1): + image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) + + grayscale = image.mean(dim=1, keepdim=True) + grayscale = grayscale.repeat(1, 3, 1, 1) + inverted = 1 - grayscale + + blur_sigma = blur_radius / 3 + blurred = self.gaussian_blur(inverted, blur_radius, blur_sigma) + + final_image = self.dodge(blurred, grayscale) + + if sharpen_alpha != 0.0: + final_image = self.sharpen(final_image, 1, sharpen_alpha) + + final_image = final_image.permute(0, 2, 3, 1) # Back to (B, H, W, C) + + return (final_image,) + + def dodge(self, front: torch.Tensor, back: torch.Tensor) -> torch.Tensor: + result = back / (1 - front + 1e-7) + result = torch.clamp(result, 0, 1) + return result + + def gaussian_blur(self, image: torch.Tensor, blur_radius: int, sigma: float): + if blur_radius == 0: + return image + + batch_size, channels, height, width = image.shape + + kernel_size = blur_radius * 2 + 1 + kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1) + + blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels) + + return blurred + + def sharpen(self, image: torch.Tensor, blur_radius: int, alpha: float): + if blur_radius == 0: + return image + + batch_size, channels, height, width = image.shape + + kernel_size = blur_radius * 2 + 1 + kernel = torch.ones((kernel_size, kernel_size), dtype=torch.float32) * -1 + center = kernel_size // 2 + kernel[center, center] = kernel_size**2 + kernel *= alpha + kernel = kernel.repeat(channels, 1, 1).unsqueeze(1) + + sharpened = F.conv2d(image, kernel, padding=center, groups=channels) + + result = torch.clamp(sharpened, 0, 1) + + return result + class PixelSort: def __init__(self): pass @@ -819,6 +893,12 @@ class Solarize: solarized_image = torch.clamp(solarized_image, 0, 1) return (solarized_image,) +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] @@ -920,6 +1000,7 @@ NODE_CLASS_MAPPINGS = { "DodgeAndBurn": DodgeAndBurn, "FilmGrain": FilmGrain, "Glow": Glow, + "PencilSketch": PencilSketch, "PixelSort": PixelSort, "Pixelize": Pixelize, "Quantize": Quantize,