slight refactoring of existing nodes

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):
+    def canny(self, image: torch.Tensor, lower_threshold: int, upper_threshold: int):
-        tensor_img = image.numpy()[0]
+        tensor_image = image.numpy()[0]
-        tensor_img = (cv2.cvtColor(tensor_img, cv2.COLOR_BGR2GRAY) * 255).astype(np.uint8)
+        gray_image = (cv2.cvtColor(tensor_image, cv2.COLOR_BGR2GRAY) * 255).astype(np.uint8)
-        canny = cv2.Canny(tensor_img, lower_threshold, upper_threshold)
+        canny = cv2.Canny(gray_image, lower_threshold, upper_threshold)
         tensor = torch.from_numpy(canny).unsqueeze(0)
         return (tensor,)
 

color_correct.py

@@ -56,8 +56,8 @@ class ColorCorrect:
 
     CATEGORY = "postprocessing"
 
-    def color_correct(self, image, temperature, hue, brightness, contrast, saturation, gamma):
+    def color_correct(self, image: torch.Tensor, temperature: float, hue: float, brightness: float, contrast: float, saturation: float, gamma: float):
-        tensor_img = image.numpy()[0]
+        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_image * 255).astype(np.uint8))
-        modified_image = Image.fromarray((tensor_img * 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)
 

dither.py

@@ -11,7 +11,7 @@ class Dither:
                 "image": ("IMAGE",),
                 "bits": ("INT", {
                     "default": 4,
-                    "min": 0,
+                    "min": 1,
                     "max": 8,
                     "step": 1
                 }),
@@ -23,36 +23,34 @@ class Dither:
 
     CATEGORY = "postprocessing"
 
-    def dither(self, image, bits):
+    def dither(self, image: torch.Tensor, bits: int):
-        tensor_img = image[0]
+        tensor_image = image.numpy()[0]
-        height, width, _ = tensor_img.shape
+        img = (tensor_image * 255)
-        out = tensor_img.clone()
+        height, width, _ = img.shape
-        levels = 2 ** bits - 1
+
+        scale = 255 / (2**bits - 1)
 
         for y in range(height):
             for x in range(width):
-                old_pixel = out[y, x].clone()
+                old_pixel = img[y, x].copy()
-                new_pixel = torch.round(old_pixel * levels) / levels
+                new_pixel = np.round(old_pixel / scale) * scale
-                out[y, x] = new_pixel
+                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)
+                    img[y, x + 1] += quant_error * 7 / 16
-                if x - 1 >= 0 and y + 1 < height:
-                    out[y + 1, x - 1] += error * 3/16
                 if y + 1 < height:
-                    out[y + 1, x] += error * 5/16
+                    if x - 1 >= 0:
-                if x + 1 < width and y + 1 < height:
+                        img[y + 1, x - 1] += quant_error * 3 / 16
-                    out[y + 1, x + 1] += error * 1/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
 }

gaussian_blur.py

@@ -30,9 +30,9 @@ class GaussianBlur:
 
     CATEGORY = "postprocessing"
 
-    def blur(self, image, kernel_size, sigma):
+    def blur(self, image: torch.Tensor, kernel_size: int, sigma: float):
-        tensor_img = image.numpy()[0]
+        tensor_image = image.numpy()[0]
-        blurred = cv2.GaussianBlur(tensor_img, (kernel_size, kernel_size), sigma)
+        blurred = cv2.GaussianBlur(tensor_image, (kernel_size, kernel_size), sigma)
         tensor = torch.from_numpy(blurred).unsqueeze(0)
         return (tensor,)
 

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,)

sharpen.py

@@ -31,17 +31,18 @@ class Sharpen:
 
     CATEGORY = "postprocessing"
 
-    def sharpen(self, image, alpha, kernel_size):
+    def sharpen(self, image: torch.Tensor, kernel_size: int, alpha: float):
-        tensor_img = image.numpy()[0]
+        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 = {