Adds Kuwahara Blur and Parabolize

post_processing/kuwahara_blur.py

@@ -0,0 +1,84 @@
+import cv2
+import numpy as np
+import multiprocessing as mp
+import torch
+
+class KuwaharaBlur:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "blur_radius": ("INT", {
+                    "default": 3,
+                    "min": 0,
+                    "max": 31,
+                    "step": 1
+                }),
+                "method": (["mean", "gaussian"],),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "apply_kuwahara_filter"
+
+    CATEGORY = "postprocessing"
+
+    def apply_kuwahara_filter(self, image: np.ndarray, blur_radius: int, method: str):
+        if blur_radius == 0:
+            return (image,)
+
+        out = torch.zeros_like(image)
+        batch_size, height, width, channels = image.shape
+
+        for b in range(batch_size):
+            image = image[b].cpu().numpy() * 255.0
+            image = image.astype(np.uint8)
+
+            out[b] = torch.from_numpy(kuwahara(image, method=method, radius=blur_radius)) / 255.0
+
+        return (out,)
+
+def kuwahara(orig_img, method="mean", radius=3, sigma=None):
+    if method == "gaussian" and sigma is None:
+        sigma = -1
+
+    image = orig_img.astype(np.float32, copy=False)
+    avgs = np.empty((4, *image.shape), dtype=image.dtype)
+    stddevs = np.empty((4, *image.shape[:2]), dtype=image.dtype)
+    image_2d = cv2.cvtColor(orig_img, cv2.COLOR_BGR2GRAY).astype(image.dtype, copy=False)
+    avgs_2d = np.empty((4, *image.shape[:2]), dtype=image.dtype)
+
+    squared_img = image_2d ** 2
+
+    if method == "mean":
+        kxy = np.ones(radius + 1, dtype=image.dtype) / (radius + 1)
+    elif method == "gaussian":
+        kxy = cv2.getGaussianKernel(2 * radius + 1, sigma, ktype=cv2.CV_32F)
+        kxy /= kxy[radius:].sum()
+        klr = np.array([kxy[:radius+1], kxy[radius:]])
+        kindexes = [[1, 1], [1, 0], [0, 1], [0, 0]]
+
+    shift = [(0, 0), (0, radius), (radius, 0), (radius, radius)]
+
+    for k in range(4):
+        if method == "mean":
+            kx, ky = kxy, kxy
+        else:
+            kx, ky = klr[kindexes[k]]
+        cv2.sepFilter2D(image, -1, kx, ky, avgs[k], shift[k])
+        cv2.sepFilter2D(image_2d, -1, kx, ky, avgs_2d[k], shift[k])
+        cv2.sepFilter2D(squared_img, -1, kx, ky, stddevs[k], shift[k])
+        stddevs[k] = stddevs[k] - avgs_2d[k] ** 2
+
+    indices = np.argmin(stddevs, axis=0)
+    filtered = np.take_along_axis(avgs, indices[None,...,None], 0).reshape(image.shape)
+
+    return filtered.astype(orig_img.dtype)
+
+NODE_CLASS_MAPPINGS = {
+    "KuwaharaBlur": KuwaharaBlur
+}

post_processing/parabolize.py

@@ -0,0 +1,45 @@
+import torch
+
+class Parabolize:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "coeff": ("FLOAT", {
+                    "default": 1.0,
+                    "min": -10.0,
+                    "max": 10.0,
+                    "step": 0.1
+                }),
+                "vertex_x": ("FLOAT", {
+                    "default": 0.5,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.1
+                }),
+                "vertex_y": ("FLOAT", {
+                    "default": 0.5,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.1
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "parabolize_image"
+
+    CATEGORY = "postprocessing"
+
+    def parabolize_image(self, image: torch.Tensor, coeff: float, vertex_x: float, vertex_y: float):
+        parabolized_image = coeff * torch.pow(image - vertex_x, 2) + vertex_y
+        parabolized_image = torch.clamp(parabolized_image, 0, 1)
+        return (parabolized_image,)
+
+NODE_CLASS_MAPPINGS = {
+    "Parabolize": Parabolize,
+}