From 8e153545126d445ddc0df694c6b701106fa8f443 Mon Sep 17 00:00:00 2001 From: EllangoK Date: Sat, 1 Apr 2023 01:00:43 -0400 Subject: [PATCH] adds pixelize effect --- README.md | 5 +- post_processing/color_correct.py | 18 ++--- post_processing/pixelize.py | 46 ++++++++++++ post_processing_nodes.py | 116 ++++++++++++++++++++++++++++--- 4 files changed, 166 insertions(+), 19 deletions(-) create mode 100644 post_processing/pixelize.py diff --git a/README.md b/README.md index 440c04b..901f687 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ A collection of post processing nodes for [ComfyUI](https://github.com/comfyanon - GaussianBlur: Applies a Gaussian blur to the input image, softening the details - KMeansQuantize: Reduce the amount of colors in an image from 0-256 - 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 - Sharpen: Enhances the details in an image by applying a sharpening filter ## Example workflow @@ -27,4 +28,6 @@ By default `post_processing_nodes.py` should have all of the combined nodes. If or just run - python combine_files.py -h for more help \ No newline at end of file + python combine_files.py -h + + for more help \ No newline at end of file diff --git a/post_processing/color_correct.py b/post_processing/color_correct.py index 8bcb309..4f1bf6a 100644 --- a/post_processing/color_correct.py +++ b/post_processing/color_correct.py @@ -61,18 +61,18 @@ class ColorCorrect: batch_size, height, width, _ = image.shape result = torch.zeros_like(image) + brightness /= 100 + contrast /= 100 + saturation /= 100 + temperature /= 100 + + brightness = 1 + brightness + contrast = 1 + contrast + saturation = 1 + saturation + for b in range(batch_size): tensor_image = image[b].numpy() - brightness /= 100 - contrast /= 100 - saturation /= 100 - temperature /= 100 - - brightness = 1 + brightness - contrast = 1 + contrast - saturation = 1 + saturation - modified_image = Image.fromarray((tensor_image * 255).astype(np.uint8)) # brightness diff --git a/post_processing/pixelize.py b/post_processing/pixelize.py new file mode 100644 index 0000000..1730678 --- /dev/null +++ b/post_processing/pixelize.py @@ -0,0 +1,46 @@ +import torch +import torch.nn.functional as F + +class Pixelize: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "pixel_size": ("INT", { + "default": 8, + "min": 2, + "max": 128, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply_pixelize" + + CATEGORY = "postprocessing" + + def apply_pixelize(self, image: torch.Tensor, pixel_size: int): + pixelized_image = self.pixelize_image(image, pixel_size) + pixelized_image = torch.clamp(pixelized_image, 0, 1) + return (pixelized_image,) + + def pixelize_image(self, image: torch.Tensor, pixel_size: int): + batch_size, height, width, channels = image.shape + new_height = height // pixel_size + new_width = width // pixel_size + + image = image.permute(0, 3, 1, 2) + image = F.avg_pool2d(image, kernel_size=pixel_size, stride=pixel_size) + image = F.interpolate(image, size=(height, width), mode='nearest') + image = image.permute(0, 2, 3, 1) + + return image + +NODE_CLASS_MAPPINGS = { + "Pixelize": Pixelize, +} diff --git a/post_processing_nodes.py b/post_processing_nodes.py index 2a50096..29bfc3f 100644 --- a/post_processing_nodes.py +++ b/post_processing_nodes.py @@ -152,18 +152,18 @@ class ColorCorrect: batch_size, height, width, _ = image.shape result = torch.zeros_like(image) + brightness /= 100 + contrast /= 100 + saturation /= 100 + temperature /= 100 + + brightness = 1 + brightness + contrast = 1 + contrast + saturation = 1 + saturation + for b in range(batch_size): tensor_image = image[b].numpy() - brightness /= 100 - contrast /= 100 - saturation /= 100 - temperature /= 100 - - brightness = 1 + brightness - contrast = 1 + contrast - saturation = 1 + saturation - modified_image = Image.fromarray((tensor_image * 255).astype(np.uint8)) # brightness @@ -449,6 +449,62 @@ class GaussianBlur: return (blurred,) +class Glow: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "intensity": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 5.0, + "step": 0.01 + }), + "blur_radius": ("INT", { + "default": 5, + "min": 1, + "max": 50, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply_glow" + + CATEGORY = "postprocessing" + + def apply_glow(self, image: torch.Tensor, intensity: float, blur_radius: int): + blurred_image = self.gaussian_blur(image, 2 * blur_radius + 1) + glowing_image = self.add_glow(image, blurred_image, intensity) + 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) + + 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) + blurred = blurred.permute(0, 2, 3, 1) + + return blurred + + def add_glow(self, img, blurred_img, intensity): + return img + blurred_img * intensity + class KMeansQuantize: def __init__(self): pass @@ -549,6 +605,46 @@ class PixelSort: return (result,) +class Pixelize: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "pixel_size": ("INT", { + "default": 8, + "min": 2, + "max": 128, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply_pixelize" + + CATEGORY = "postprocessing" + + def apply_pixelize(self, image: torch.Tensor, pixel_size: int): + pixelized_image = self.pixelize_image(image, pixel_size) + pixelized_image = torch.clamp(pixelized_image, 0, 1) + return (pixelized_image,) + + def pixelize_image(self, image: torch.Tensor, pixel_size: int): + batch_size, height, width, channels = image.shape + new_height = height // pixel_size + new_width = width // pixel_size + + image = image.permute(0, 3, 1, 2) + image = F.avg_pool2d(image, kernel_size=pixel_size, stride=pixel_size) + image = F.interpolate(image, size=(height, width), mode='nearest') + image = image.permute(0, 2, 3, 1) + + return image + class Sharpen: def __init__(self): pass @@ -693,7 +789,9 @@ NODE_CLASS_MAPPINGS = { "Dither": Dither, "FilmGrain": FilmGrain, "GaussianBlur": GaussianBlur, + "Glow": Glow, "KMeansQuantize": KMeansQuantize, "PixelSort": PixelSort, + "Pixelize": Pixelize, "Sharpen": Sharpen, }