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

adds PencilSketch effect, mimics a drawn sketch

Browse Source 
also extracts out gaussian_kernel func
This commit is contained in:
EllangoK
2023-04-03 17:01:00 -04:00
parent 99956683a0
commit 038a36d25e
5 changed files with 209 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -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
post_processing/blur.py
+7 -7
View File
@@ -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()
post_processing/glow.py
+8 -8
View File
@@ -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()
post_processing/pencil_sketch.py
+97
View File
@@ -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()
post_processing_nodes.py
+96 -15
View File
@@ -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,