84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
import cv2
|
|
import numpy as np
|
|
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/Filters"
|
|
|
|
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
|
|
}
|