386 lines
11 KiB
Python
386 lines
11 KiB
Python
import torch
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
import cv2
|
|
from cv2.ximgproc import guidedFilter
|
|
import copy
|
|
|
|
class AlphaClean:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"radius": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1
|
|
}),
|
|
"fill_holes": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 16,
|
|
"step": 1
|
|
}),
|
|
"white_threshold": ("FLOAT", {
|
|
"default": 0.9,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"extra_clip": ("FLOAT", {
|
|
"default": 0.98,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "alpha_clean"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def alpha_clean(self, images: torch.Tensor, radius: int, fill_holes: int, white_threshold: float, extra_clip: float):
|
|
|
|
d = radius * 2 + 1
|
|
i_dup = copy.deepcopy(images.cpu().numpy())
|
|
|
|
for index, image in enumerate(i_dup):
|
|
|
|
cleaned = cv2.bilateralFilter(image, 9, 0.05, 8)
|
|
|
|
alpha = np.clip((image - white_threshold) / (1 - white_threshold), 0, 1)
|
|
rgb = image * alpha
|
|
|
|
alpha = cv2.GaussianBlur(alpha, (d,d), 0) * 0.99 + np.average(alpha) * 0.01
|
|
rgb = cv2.GaussianBlur(rgb, (d,d), 0) * 0.99 + np.average(rgb) * 0.01
|
|
|
|
rgb = rgb / np.clip(alpha, 0.00001, 1)
|
|
rgb = rgb * extra_clip
|
|
|
|
cleaned = np.clip(cleaned / rgb, 0, 1)
|
|
|
|
if fill_holes > 0:
|
|
fD = fill_holes * 2 + 1
|
|
gamma = cleaned * cleaned
|
|
kD = np.ones((fD, fD), np.uint8)
|
|
kE = np.ones((fD + 2, fD + 2), np.uint8)
|
|
gamma = cv2.dilate(gamma, kD, iterations=1)
|
|
gamma = cv2.erode(gamma, kE, iterations=1)
|
|
gamma = cv2.GaussianBlur(gamma, (fD, fD), 0)
|
|
cleaned = np.maximum(cleaned, gamma)
|
|
|
|
i_dup[index] = cleaned
|
|
|
|
return (torch.from_numpy(i_dup),)
|
|
|
|
class BlurImageFast:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"radius_x": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
"radius_y": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "blur_image"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def blur_image(self, images, radius_x, radius_y):
|
|
|
|
if radius_x + radius_y == 0:
|
|
return (images,)
|
|
|
|
dx = radius_x * 2 + 1
|
|
dy = radius_y * 2 + 1
|
|
|
|
dup = copy.deepcopy(images.cpu().numpy())
|
|
|
|
for index, image in enumerate(dup):
|
|
dup[index] = cv2.GaussianBlur(image, (dx, dy), 0)
|
|
|
|
return (torch.from_numpy(dup),)
|
|
|
|
class BlurMaskFast:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"masks": ("MASK",),
|
|
"radius_x": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
"radius_y": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "blur_mask"
|
|
|
|
CATEGORY = "mask/filters"
|
|
|
|
def blur_mask(self, masks, radius_x, radius_y):
|
|
|
|
if radius_x + radius_y == 0:
|
|
return (masks,)
|
|
|
|
dx = radius_x * 2 + 1
|
|
dy = radius_y * 2 + 1
|
|
|
|
dup = copy.deepcopy(masks.cpu().numpy())
|
|
|
|
for index, mask in enumerate(dup):
|
|
dup[index] = cv2.GaussianBlur(mask, (dx, dy), 0)
|
|
|
|
return (torch.from_numpy(dup),)
|
|
|
|
class DilateErodeMask:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"masks": ("MASK",),
|
|
"radius": ("INT", {
|
|
"default": 0,
|
|
"min": -1023,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
"shape": (["box", "circle"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "dilate_mask"
|
|
|
|
CATEGORY = "mask/filters"
|
|
|
|
def dilate_mask(self, masks, radius, shape):
|
|
|
|
if radius == 0:
|
|
return (masks,)
|
|
|
|
s = abs(radius)
|
|
d = s * 2 + 1
|
|
k = np.zeros((d, d), np.uint8)
|
|
if shape == "circle":
|
|
k = cv2.circle(k, (s,s), s, 1, -1)
|
|
else:
|
|
k += 1
|
|
|
|
dup = copy.deepcopy(masks.cpu().numpy())
|
|
|
|
for index, mask in enumerate(dup):
|
|
if radius > 0:
|
|
dup[index] = cv2.dilate(mask, k, iterations=1)
|
|
else:
|
|
dup[index] = cv2.erode(mask, k, iterations=1)
|
|
|
|
return (torch.from_numpy(dup),)
|
|
|
|
class EnhanceDetail:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"filter_radius": ("INT", {
|
|
"default": 2,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1
|
|
}),
|
|
"sigma": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.01,
|
|
"max": 100.0,
|
|
"step": 0.01
|
|
}),
|
|
"denoise": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.01
|
|
}),
|
|
"detail_mult": ("FLOAT", {
|
|
"default": 2.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "enhance"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def enhance(self, images: torch.Tensor, filter_radius: int, sigma: float, denoise: float, detail_mult: float):
|
|
|
|
if filter_radius == 0:
|
|
return (images,)
|
|
|
|
d = filter_radius * 2 + 1
|
|
s = sigma / 10
|
|
n = denoise / 10
|
|
|
|
dup = copy.deepcopy(images.cpu().numpy())
|
|
|
|
for index, image in enumerate(dup):
|
|
imgB = image
|
|
if denoise>0.0:
|
|
imgB = cv2.bilateralFilter(image, d, n, d)
|
|
|
|
imgG = np.clip(guidedFilter(image, image, d, s), 0.001, 1)
|
|
|
|
details = (imgB/imgG - 1) * detail_mult + 1
|
|
dup[index] = np.clip(details*imgG - imgB + image, 0, 1)
|
|
|
|
return (torch.from_numpy(dup),)
|
|
|
|
class GuidedFilterAlpha:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"alpha": ("IMAGE",),
|
|
"filter_radius": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1
|
|
}),
|
|
"sigma": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "guided_filter_alpha"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def guided_filter_alpha(self, images: torch.Tensor, alpha: torch.Tensor, filter_radius: int, sigma: float):
|
|
|
|
d = filter_radius * 2 + 1
|
|
s = sigma / 10
|
|
|
|
i_dup = copy.deepcopy(images.cpu().numpy())
|
|
a_dup = copy.deepcopy(alpha.cpu().numpy())
|
|
|
|
for index, image in enumerate(i_dup):
|
|
alpha_work = a_dup[index]
|
|
i_dup[index] = guidedFilter(image, alpha_work, d, s)
|
|
|
|
return (torch.from_numpy(i_dup),)
|
|
|
|
class RemapRange:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"blackpoint": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"whitepoint": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "remap"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def remap(self, image: torch.Tensor, blackpoint: float, whitepoint: float):
|
|
|
|
bp = min(blackpoint, whitepoint - 0.001)
|
|
scale = 1 / (whitepoint - bp)
|
|
|
|
i_dup = copy.deepcopy(image.cpu().numpy())
|
|
i_dup = np.clip((i_dup - bp) * scale, 0.0, 1.0)
|
|
|
|
return (torch.from_numpy(i_dup),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AlphaClean": AlphaClean,
|
|
"BlurImageFast": BlurImageFast,
|
|
"BlurMaskFast": BlurMaskFast,
|
|
"DilateErodeMask": DilateErodeMask,
|
|
"EnhanceDetail": EnhanceDetail,
|
|
"GuidedFilterAlpha": GuidedFilterAlpha,
|
|
"RemapRange": RemapRange,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AlphaClean": "Alpha Clean",
|
|
"BlurImageFast": "Blur Image (Fast)",
|
|
"BlurMaskFast": "Blur Mask (Fast)",
|
|
"DilateErodeMask": "Dilate/Erode Mask",
|
|
"EnhanceDetail": "Enhance Detail",
|
|
"GuidedFilterAlpha": "Guided Filter Alpha",
|
|
"RemapRange": "Remap Range",
|
|
} |