Skip to content

Commit 038a36d

Browse files
EllangoKEllangoK
committed
adds PencilSketch effect, mimics a drawn sketch
also extracts out gaussian_kernel func
1 parent 9995668 commit 038a36d

File tree

5 files changed

+209
-30
lines changed

5 files changed

+209
-30
lines changed

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ A collection of post processing nodes for [ComfyUI](https://github.com/comfyanon
1313
- DodgeAndBurn: Adjusts image brightness using dodge and burn effects based on a mask and intensity.
1414
- FilmGrain: Adds a film grain effect to the image, along with options to control the temperature, and vignetting
1515
- Glow: Applies a blur with a specified radius and then blends it with the original image. Creates a nice glowing effect.
16+
- PencilSketch: Converts an image into a hand-drawn pencil sketch style.
1617
- PixelSort: Rearranges the pixels in the input image based on their values, and input mask. Creates a cool glitch like effect.
1718
- Pixelize: Applies a pixelization effect, simulating the reducing of resolution
1819
- Quantize: Set and dither the amount of colors in an image from 0-256, reducing color information

‎post_processing/blur.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,20 +30,14 @@ def INPUT_TYPES(s):
3030

3131
CATEGORY = "postprocessing"
3232

33-
def gaussian_kernel(self, kernel_size: int, sigma: float):
34-
x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij")
35-
d = torch.sqrt(x * x + y * y)
36-
g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
37-
return g / g.sum()
38-
3933
def blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
4034
if blur_radius == 0:
4135
return (image,)
4236

4337
batch_size, height, width, channels = image.shape
4438

4539
kernel_size = blur_radius * 2 + 1
46-
kernel = self.gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
40+
kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
4741

4842
image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
4943
blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels)
@@ -54,3 +48,9 @@ def blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
5448
NODE_CLASS_MAPPINGS = {
5549
"Blur": Blur
5650
}
51+
52+
def gaussian_kernel(kernel_size: int, sigma: float):
53+
x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij")
54+
d = torch.sqrt(x * x + y * y)
55+
g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
56+
return g / g.sum()

‎post_processing/glow.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,17 +36,11 @@ def apply_glow(self, image: torch.Tensor, intensity: float, blur_radius: int):
3636
glowing_image = torch.clamp(glowing_image, 0, 1)
3737
return (glowing_image,)
3838

39-
def gaussian_kernel(self, kernel_size: int):
40-
sigma = (kernel_size - 1) / 6
41-
x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size))
42-
d = torch.sqrt(x * x + y * y)
43-
g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
44-
return g / g.sum()
45-
4639
def gaussian_blur(self, image: torch.Tensor, kernel_size: int):
4740
batch_size, height, width, channels = image.shape
4841

49-
kernel = self.gaussian_kernel(kernel_size).repeat(channels, 1, 1).unsqueeze(1)
42+
sigma = (kernel_size - 1) / 6
43+
kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
5044

5145
image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
5246
blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels)
@@ -60,3 +54,9 @@ def add_glow(self, img, blurred_img, intensity):
6054
NODE_CLASS_MAPPINGS = {
6155
"Glow": Glow,
6256
}
57+
58+
def gaussian_kernel(kernel_size: int, sigma: float):
59+
x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij")
60+
d = torch.sqrt(x * x + y * y)
61+
g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
62+
return g / g.sum()

‎post_processing/pencil_sketch.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
import torch
2+
import torch.nn.functional as F
3+
4+
class PencilSketch:
5+
def __init__(self):
6+
pass
7+
8+
@classmethod
9+
def INPUT_TYPES(cls):
10+
return {
11+
"required": {
12+
"image": ("IMAGE",),
13+
"blur_radius": ("INT", {
14+
"default": 5,
15+
"min": 1,
16+
"max": 31,
17+
"step": 1
18+
}),
19+
"sharpen_alpha": ("FLOAT", {
20+
"default": 1.0,
21+
"min": 0.0,
22+
"max": 10.0,
23+
"step": 0.1
24+
}),
25+
},
26+
}
27+
28+
RETURN_TYPES = ("IMAGE",)
29+
FUNCTION = "apply_sketch"
30+
31+
CATEGORY = "postprocessing"
32+
33+
def apply_sketch(self, image: torch.Tensor, blur_radius: int = 5, sharpen_alpha: float = 1):
34+
image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
35+
36+
grayscale = image.mean(dim=1, keepdim=True)
37+
grayscale = grayscale.repeat(1, 3, 1, 1)
38+
inverted = 1 - grayscale
39+
40+
blur_sigma = blur_radius / 3
41+
blurred = self.gaussian_blur(inverted, blur_radius, blur_sigma)
42+
43+
final_image = self.dodge(blurred, grayscale)
44+
45+
if sharpen_alpha != 0.0:
46+
final_image = self.sharpen(final_image, 1, sharpen_alpha)
47+
48+
final_image = final_image.permute(0, 2, 3, 1) # Back to (B, H, W, C)
49+
50+
return (final_image,)
51+
52+
def dodge(self, front: torch.Tensor, back: torch.Tensor) -> torch.Tensor:
53+
result = back / (1 - front + 1e-7)
54+
result = torch.clamp(result, 0, 1)
55+
return result
56+
57+
def gaussian_blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
58+
if blur_radius == 0:
59+
return image
60+
61+
batch_size, channels, height, width = image.shape
62+
63+
kernel_size = blur_radius * 2 + 1
64+
kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
65+
66+
blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels)
67+
68+
return blurred
69+
70+
def sharpen(self, image: torch.Tensor, blur_radius: int, alpha: float):
71+
if blur_radius == 0:
72+
return image
73+
74+
batch_size, channels, height, width = image.shape
75+
76+
kernel_size = blur_radius * 2 + 1
77+
kernel = torch.ones((kernel_size, kernel_size), dtype=torch.float32) * -1
78+
center = kernel_size // 2
79+
kernel[center, center] = kernel_size**2
80+
kernel *= alpha
81+
kernel = kernel.repeat(channels, 1, 1).unsqueeze(1)
82+
83+
sharpened = F.conv2d(image, kernel, padding=center, groups=channels)
84+
85+
result = torch.clamp(sharpened, 0, 1)
86+
87+
return result
88+
89+
NODE_CLASS_MAPPINGS = {
90+
"PencilSketch": PencilSketch,
91+
}
92+
93+
def gaussian_kernel(kernel_size: int, sigma: float):
94+
x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij")
95+
d = torch.sqrt(x * x + y * y)
96+
g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
97+
return g / g.sum()

‎post_processing_nodes.py‎

Lines changed: 96 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import cv2
44
import numpy as np
55
from PIL import Image, ImageEnhance
6+
from PIL import Image
67

78

89
class ArithmeticBlend:
@@ -149,20 +150,14 @@ def INPUT_TYPES(s):
149150

150151
CATEGORY = "postprocessing"
151152

152-
def gaussian_kernel(self, kernel_size: int, sigma: float):
153-
x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij")
154-
d = torch.sqrt(x * x + y * y)
155-
g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
156-
return g / g.sum()
157-
158153
def blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
159154
if blur_radius == 0:
160155
return (image,)
161156

162157
batch_size, height, width, channels = image.shape
163158

164159
kernel_size = blur_radius * 2 + 1
165-
kernel = self.gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
160+
kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
166161

167162
image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
168163
blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels)
@@ -593,17 +588,11 @@ def apply_glow(self, image: torch.Tensor, intensity: float, blur_radius: int):
593588
glowing_image = torch.clamp(glowing_image, 0, 1)
594589
return (glowing_image,)
595590

596-
def gaussian_kernel(self, kernel_size: int):
597-
sigma = (kernel_size - 1) / 6
598-
x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size))
599-
d = torch.sqrt(x * x + y * y)
600-
g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
601-
return g / g.sum()
602-
603591
def gaussian_blur(self, image: torch.Tensor, kernel_size: int):
604592
batch_size, height, width, channels = image.shape
605593

606-
kernel = self.gaussian_kernel(kernel_size).repeat(channels, 1, 1).unsqueeze(1)
594+
sigma = (kernel_size - 1) / 6
595+
kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
607596

608597
image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
609598
blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels)
@@ -614,6 +603,91 @@ def gaussian_blur(self, image: torch.Tensor, kernel_size: int):
614603
def add_glow(self, img, blurred_img, intensity):
615604
return img + blurred_img * intensity
616605

606+
class PencilSketch:
607+
def __init__(self):
608+
pass
609+
610+
@classmethod
611+
def INPUT_TYPES(cls):
612+
return {
613+
"required": {
614+
"image": ("IMAGE",),
615+
"blur_radius": ("INT", {
616+
"default": 5,
617+
"min": 1,
618+
"max": 31,
619+
"step": 1
620+
}),
621+
"sharpen_alpha": ("FLOAT", {
622+
"default": 1.0,
623+
"min": 0.0,
624+
"max": 10.0,
625+
"step": 0.1
626+
}),
627+
},
628+
}
629+
630+
RETURN_TYPES = ("IMAGE",)
631+
FUNCTION = "apply_sketch"
632+
633+
CATEGORY = "postprocessing"
634+
635+
def apply_sketch(self, image: torch.Tensor, blur_radius: int = 5, sharpen_alpha: float = 1):
636+
image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
637+
638+
grayscale = image.mean(dim=1, keepdim=True)
639+
grayscale = grayscale.repeat(1, 3, 1, 1)
640+
inverted = 1 - grayscale
641+
642+
blur_sigma = blur_radius / 3
643+
blurred = self.gaussian_blur(inverted, blur_radius, blur_sigma)
644+
645+
final_image = self.dodge(blurred, grayscale)
646+
647+
if sharpen_alpha != 0.0:
648+
final_image = self.sharpen(final_image, 1, sharpen_alpha)
649+
650+
final_image = final_image.permute(0, 2, 3, 1) # Back to (B, H, W, C)
651+
652+
return (final_image,)
653+
654+
def dodge(self, front: torch.Tensor, back: torch.Tensor) -> torch.Tensor:
655+
result = back / (1 - front + 1e-7)
656+
result = torch.clamp(result, 0, 1)
657+
return result
658+
659+
def gaussian_blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
660+
if blur_radius == 0:
661+
return image
662+
663+
batch_size, channels, height, width = image.shape
664+
665+
kernel_size = blur_radius * 2 + 1
666+
kernel = gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1)
667+
668+
blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels)
669+
670+
return blurred
671+
672+
def sharpen(self, image: torch.Tensor, blur_radius: int, alpha: float):
673+
if blur_radius == 0:
674+
return image
675+
676+
batch_size, channels, height, width = image.shape
677+
678+
kernel_size = blur_radius * 2 + 1
679+
kernel = torch.ones((kernel_size, kernel_size), dtype=torch.float32) * -1
680+
center = kernel_size // 2
681+
kernel[center, center] = kernel_size**2
682+
kernel *= alpha
683+
kernel = kernel.repeat(channels, 1, 1).unsqueeze(1)
684+
685+
sharpened = F.conv2d(image, kernel, padding=center, groups=channels)
686+
687+
result = torch.clamp(sharpened, 0, 1)
688+
689+
return result
690+
617691
class PixelSort:
618692
def __init__(self):
619693
pass
@@ -819,6 +893,12 @@ def solarize_image(self, image: torch.Tensor, threshold: float):
819893
solarized_image = torch.clamp(solarized_image, 0, 1)
820894
return (solarized_image,)
821895

896+
def gaussian_kernel(kernel_size: int, sigma: float):
897+
x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij")
898+
d = torch.sqrt(x * x + y * y)
899+
g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
900+
return g / g.sum()
901+
822902
def sort_span(span, sort_by, reverse_sorting):
823903
if sort_by == 'H':
824904
key = lambda x: x[1][0]
@@ -920,6 +1000,7 @@ def pixel_sort(img, mask, horizontal_sort=False, span_limit=None, sort_by='H', r
9201000
"DodgeAndBurn": DodgeAndBurn,
9211001
"FilmGrain": FilmGrain,
9221002
"Glow": Glow,
1003+
"PencilSketch": PencilSketch,
9231004
"PixelSort": PixelSort,
9241005
"Pixelize": Pixelize,
9251006
"Quantize": Quantize,

0 commit comments

Comments
 (0)