From ece18e5f0d50592af03a5efc1ff422d6194bea43 Mon Sep 17 00:00:00 2001 From: EllangoK Date: Thu, 30 Mar 2023 16:54:27 -0400 Subject: [PATCH] slight refactoring of existing nodes --- canny_edge_detect.py | 12 +++++----- color_correct.py | 9 +++----- dither.py | 46 ++++++++++++++++++-------------------- gaussian_blur.py | 8 +++---- kmeans_quantize.py | 53 ++++++++++++++++++++++++++++++++++++++++++++ sharpen.py | 7 +++--- 6 files changed, 92 insertions(+), 43 deletions(-) create mode 100644 kmeans_quantize.py diff --git a/canny_edge_detect.py b/canny_edge_detect.py index c77f8f6..2484d50 100644 --- a/canny_edge_detect.py +++ b/canny_edge_detect.py @@ -15,13 +15,13 @@ class CannyEdgeDetection: "default": 100, "min": 0, "max": 500, - "step": 1 + "step": 10 }), "upper_threshold": ("INT", { "default": 200, "min": 0, "max": 500, - "step": 1 + "step": 10 }), }, } @@ -31,10 +31,10 @@ class CannyEdgeDetection: CATEGORY = "postprocessing" - def canny(self, image, lower_threshold, upper_threshold): - tensor_img = image.numpy()[0] - tensor_img = (cv2.cvtColor(tensor_img, cv2.COLOR_BGR2GRAY) * 255).astype(np.uint8) - canny = cv2.Canny(tensor_img, lower_threshold, upper_threshold) + def canny(self, image: torch.Tensor, lower_threshold: int, upper_threshold: int): + tensor_image = image.numpy()[0] + gray_image = (cv2.cvtColor(tensor_image, cv2.COLOR_BGR2GRAY) * 255).astype(np.uint8) + canny = cv2.Canny(gray_image, lower_threshold, upper_threshold) tensor = torch.from_numpy(canny).unsqueeze(0) return (tensor,) diff --git a/color_correct.py b/color_correct.py index 68b7ffc..6e75872 100644 --- a/color_correct.py +++ b/color_correct.py @@ -56,8 +56,8 @@ class ColorCorrect: CATEGORY = "postprocessing" - def color_correct(self, image, temperature, hue, brightness, contrast, saturation, gamma): - tensor_img = image.numpy()[0] + def color_correct(self, image: torch.Tensor, temperature: float, hue: float, brightness: float, contrast: float, saturation: float, gamma: float): + tensor_image = image.numpy()[0] brightness /= 100 contrast /= 100 @@ -68,8 +68,7 @@ class ColorCorrect: contrast = 1 + contrast saturation = 1 + saturation - - modified_image = Image.fromarray((tensor_img * 255).astype(np.uint8)) + modified_image = Image.fromarray((tensor_image * 255).astype(np.uint8)) # brightness modified_image = ImageEnhance.Brightness(modified_image).enhance(brightness) @@ -97,11 +96,9 @@ class ColorCorrect: # hue hsv_img = cv2.cvtColor(modified_image, cv2.COLOR_RGB2HSV) hsv_img[:, :, 0] = (hsv_img[:, :, 0] + hue) % 360 - modified_image = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2RGB) modified_image = modified_image.astype(np.uint8) - modified_image = modified_image / 255 modified_image = torch.from_numpy(modified_image).unsqueeze(0) diff --git a/dither.py b/dither.py index 846bdad..d94ea0b 100644 --- a/dither.py +++ b/dither.py @@ -3,15 +3,15 @@ import torch class Dither: def __init__(self): pass - + @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "bits": ("INT", { - "default": 4, - "min": 0, + "default": 4, + "min": 1, "max": 8, "step": 1 }), @@ -23,36 +23,34 @@ class Dither: CATEGORY = "postprocessing" - def dither(self, image, bits): - tensor_img = image[0] - height, width, _ = tensor_img.shape - out = tensor_img.clone() - levels = 2 ** bits - 1 + def dither(self, image: torch.Tensor, bits: int): + tensor_image = image.numpy()[0] + img = (tensor_image * 255) + height, width, _ = img.shape + + scale = 255 / (2**bits - 1) for y in range(height): for x in range(width): - old_pixel = out[y, x].clone() - new_pixel = torch.round(old_pixel * levels) / levels - out[y, x] = new_pixel + old_pixel = img[y, x].copy() + new_pixel = np.round(old_pixel / scale) * scale + img[y, x] = new_pixel - error = old_pixel - new_pixel + quant_error = old_pixel - new_pixel if x + 1 < width: - out[y, x + 1] += error * (7 / 16) - if x - 1 >= 0 and y + 1 < height: - out[y + 1, x - 1] += error * 3/16 + img[y, x + 1] += quant_error * 7 / 16 if y + 1 < height: - out[y + 1, x] += error * 5/16 - if x + 1 < width and y + 1 < height: - out[y + 1, x + 1] += error * 1/16 + if x - 1 >= 0: + img[y + 1, x - 1] += quant_error * 3 / 16 + img[y + 1, x] += quant_error * 5 / 16 + if x + 1 < width: + img[y + 1, x + 1] += quant_error * 1 / 16 - out = torch.clamp(out, 0, 1).unsqueeze(0) + dithered = img / 255 + tensor = torch.from_numpy(dithered).unsqueeze(0) + return (tensor,) - return (out,) - - -# A dictionary that contains all nodes you want to export with their names -# NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "Dither": Dither } diff --git a/gaussian_blur.py b/gaussian_blur.py index c6d8b4e..db2fdd0 100644 --- a/gaussian_blur.py +++ b/gaussian_blur.py @@ -30,12 +30,12 @@ class GaussianBlur: CATEGORY = "postprocessing" - def blur(self, image, kernel_size, sigma): - tensor_img = image.numpy()[0] - blurred = cv2.GaussianBlur(tensor_img, (kernel_size, kernel_size), sigma) + def blur(self, image: torch.Tensor, kernel_size: int, sigma: float): + tensor_image = image.numpy()[0] + blurred = cv2.GaussianBlur(tensor_image, (kernel_size, kernel_size), sigma) tensor = torch.from_numpy(blurred).unsqueeze(0) return (tensor,) - + NODE_CLASS_MAPPINGS = { "GaussianBlur": GaussianBlur } \ No newline at end of file diff --git a/kmeans_quantize.py b/kmeans_quantize.py new file mode 100644 index 0000000..a8e30ed --- /dev/null +++ b/kmeans_quantize.py @@ -0,0 +1,53 @@ +import numpy as np +import cv2 +import torch + +class KMeansQuantize: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "colors": ("INT", { + "default": 16, + "min": 1, + "max": 256, + "step": 1 + }), + "precision": ("INT", { + "default": 10, + "min": 1, + "max": 100, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "kmeans_quantize" + + CATEGORY = "postprocessing" + + def kmeans_quantize(self, image: torch.Tensor, colors: int, precision: int): + tensor_image = image.numpy()[0].astype(np.float32) + img = tensor_image + + height, width, c = img.shape + + criteria = ( + cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, + precision * 5, 0.01 + ) + + img_copy = img.reshape(-1, c) + _, label, center = cv2.kmeans( + img_copy, colors, None, + criteria, 1, cv2.KMEANS_PP_CENTERS + ) + + result = center[label.flatten()].reshape(*img.shape) + tensor = torch.from_numpy(result).unsqueeze(0) + return (tensor,) diff --git a/sharpen.py b/sharpen.py index 2399a34..3266032 100644 --- a/sharpen.py +++ b/sharpen.py @@ -31,17 +31,18 @@ class Sharpen: CATEGORY = "postprocessing" - def sharpen(self, image, alpha, kernel_size): - tensor_img = image.numpy()[0] + def sharpen(self, image: torch.Tensor, kernel_size: int, alpha: float): + tensor_image = image.numpy()[0] kernel = np.ones((kernel_size, kernel_size), dtype=np.float32) * -1 center = kernel_size // 2 kernel[center, center] = kernel_size**2 kernel *= alpha - sharpened = cv2.filter2D(tensor_img, -1, kernel) + sharpened = cv2.filter2D(tensor_image, -1, kernel) tensor = torch.from_numpy(sharpened).unsqueeze(0) + tensor = torch.clamp(tensor, 0, 1) return (tensor,) NODE_CLASS_MAPPINGS = {