399 lines
14 KiB
Python
399 lines
14 KiB
Python
"""
|
|
Filter Utilities for TextureAlchemy
|
|
Denoise, edge detection, enhancement filters
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class DenoiseFilter:
|
|
"""
|
|
Noise reduction with detail preservation
|
|
Multiple denoising algorithms
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"strength": ("FLOAT", {
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"method": (["bilateral", "non_local_means", "gaussian", "median"],),
|
|
"preserve_edges": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("denoised",)
|
|
FUNCTION = "denoise"
|
|
CATEGORY = "Texture Alchemist/Filters"
|
|
|
|
def denoise(self, image, strength, method, preserve_edges):
|
|
"""Apply denoising"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Denoise Filter - {method.upper()}")
|
|
print(f"{'='*60}")
|
|
|
|
if method == "bilateral":
|
|
result = self._bilateral_filter(image, strength, preserve_edges)
|
|
elif method == "non_local_means":
|
|
result = self._nlm_filter(image, strength)
|
|
elif method == "gaussian":
|
|
result = self._gaussian_denoise(image, strength)
|
|
else: # median
|
|
result = self._median_filter(image, strength)
|
|
|
|
# Blend with original based on strength
|
|
result = image * (1.0 - strength) + result * strength
|
|
result = torch.clamp(result, 0.0, 1.0)
|
|
|
|
print(f"✓ Denoising applied")
|
|
print(f" Method: {method}")
|
|
print(f" Strength: {strength:.2f}")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (result,)
|
|
|
|
def _bilateral_filter(self, image, strength, preserve_edges):
|
|
"""Bilateral filter (edge-preserving blur)"""
|
|
sigma_space = 5.0 * strength
|
|
sigma_color = 0.1
|
|
|
|
kernel_size = max(3, int(sigma_space * 2) + 1)
|
|
if kernel_size % 2 == 0:
|
|
kernel_size += 1
|
|
|
|
# Simple bilateral approximation
|
|
gaussian = self._gaussian_denoise(image, strength)
|
|
|
|
if preserve_edges:
|
|
# Detect edges
|
|
gray = image.mean(dim=3, keepdim=True)
|
|
dx = torch.abs(gray[:, :, 1:, :] - gray[:, :, :-1, :])
|
|
dy = torch.abs(gray[:, 1:, :, :] - gray[:, :-1, :, :])
|
|
|
|
dx = F.pad(dx, (0, 0, 0, 1), mode='replicate')
|
|
dy = F.pad(dy, (0, 0, 0, 0, 0, 1), mode='replicate')
|
|
|
|
edges = (dx + dy).repeat(1, 1, 1, image.shape[3])
|
|
|
|
# Preserve edges
|
|
blend = torch.exp(-edges / sigma_color)
|
|
return image * (1.0 - blend) + gaussian * blend
|
|
|
|
return gaussian
|
|
|
|
def _nlm_filter(self, image, strength):
|
|
"""Non-local means (simplified)"""
|
|
# Simplified patch-based denoising
|
|
# Use Gaussian blur as approximation for performance
|
|
return self._gaussian_denoise(image, strength * 1.5)
|
|
|
|
def _gaussian_denoise(self, image, strength):
|
|
"""Gaussian blur denoising"""
|
|
sigma = 2.0 * strength
|
|
kernel_size = max(3, int(sigma * 4) + 1)
|
|
if kernel_size % 2 == 0:
|
|
kernel_size += 1
|
|
kernel_size = min(kernel_size, 25)
|
|
|
|
x = torch.arange(kernel_size, dtype=image.dtype, device=image.device) - kernel_size // 2
|
|
kernel_1d = torch.exp(-x ** 2 / (2 * sigma ** 2))
|
|
kernel_1d = kernel_1d / kernel_1d.sum()
|
|
|
|
image_t = image.permute(0, 3, 1, 2)
|
|
|
|
# Horizontal
|
|
kernel_h = kernel_1d.view(1, 1, 1, -1).repeat(image_t.shape[1], 1, 1, 1)
|
|
blurred = F.conv2d(F.pad(image_t, (kernel_size//2, kernel_size//2, 0, 0), mode='replicate'),
|
|
kernel_h, groups=image_t.shape[1])
|
|
|
|
# Vertical
|
|
kernel_v = kernel_1d.view(1, 1, -1, 1).repeat(blurred.shape[1], 1, 1, 1)
|
|
blurred = F.conv2d(F.pad(blurred, (0, 0, kernel_size//2, kernel_size//2), mode='replicate'),
|
|
kernel_v, groups=blurred.shape[1])
|
|
|
|
return blurred.permute(0, 2, 3, 1)
|
|
|
|
def _median_filter(self, image, strength):
|
|
"""Median filter (approximated with average pooling)"""
|
|
kernel_size = max(3, int(strength * 10) + 1)
|
|
if kernel_size % 2 == 0:
|
|
kernel_size += 1
|
|
|
|
image_t = image.permute(0, 3, 1, 2)
|
|
|
|
filtered = F.avg_pool2d(
|
|
F.pad(image_t, (kernel_size//2, kernel_size//2, kernel_size//2, kernel_size//2), mode='replicate'),
|
|
kernel_size=kernel_size,
|
|
stride=1
|
|
)
|
|
|
|
return filtered.permute(0, 2, 3, 1)
|
|
|
|
|
|
class EdgeDetection:
|
|
"""
|
|
Detect edges in images
|
|
Multiple edge detection algorithms
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"method": (["sobel", "scharr", "prewitt", "canny_approx", "laplacian"],),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.1,
|
|
"display": "number"
|
|
}),
|
|
"threshold": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Edge threshold (0=all, 1=strong only)"
|
|
}),
|
|
"invert": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("edges",)
|
|
FUNCTION = "detect"
|
|
CATEGORY = "Texture Alchemist/Filters"
|
|
|
|
def detect(self, image, method, strength, threshold, invert):
|
|
"""Detect edges"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Edge Detection - {method.upper()}")
|
|
print(f"{'='*60}")
|
|
|
|
# Convert to grayscale
|
|
gray = 0.299 * image[:, :, :, 0:1] + 0.587 * image[:, :, :, 1:2] + 0.114 * image[:, :, :, 2:3]
|
|
gray_t = gray.permute(0, 3, 1, 2)
|
|
|
|
if method == "sobel":
|
|
sobel_x = torch.tensor([[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]],
|
|
dtype=gray.dtype, device=gray.device).view(1, 1, 3, 3)
|
|
sobel_y = torch.tensor([[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]],
|
|
dtype=gray.dtype, device=gray.device).view(1, 1, 3, 3)
|
|
|
|
dx = F.conv2d(F.pad(gray_t, (1, 1, 1, 1), mode='replicate'), sobel_x)
|
|
dy = F.conv2d(F.pad(gray_t, (1, 1, 1, 1), mode='replicate'), sobel_y)
|
|
|
|
elif method == "scharr":
|
|
scharr_x = torch.tensor([[[-3, 0, 3], [-10, 0, 10], [-3, 0, 3]]],
|
|
dtype=gray.dtype, device=gray.device).view(1, 1, 3, 3)
|
|
scharr_y = torch.tensor([[[-3, -10, -3], [0, 0, 0], [3, 10, 3]]],
|
|
dtype=gray.dtype, device=gray.device).view(1, 1, 3, 3)
|
|
|
|
dx = F.conv2d(F.pad(gray_t, (1, 1, 1, 1), mode='replicate'), scharr_x)
|
|
dy = F.conv2d(F.pad(gray_t, (1, 1, 1, 1), mode='replicate'), scharr_y)
|
|
|
|
elif method == "prewitt":
|
|
prewitt_x = torch.tensor([[[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]]],
|
|
dtype=gray.dtype, device=gray.device).view(1, 1, 3, 3)
|
|
prewitt_y = torch.tensor([[[-1, -1, -1], [0, 0, 0], [1, 1, 1]]],
|
|
dtype=gray.dtype, device=gray.device).view(1, 1, 3, 3)
|
|
|
|
dx = F.conv2d(F.pad(gray_t, (1, 1, 1, 1), mode='replicate'), prewitt_x)
|
|
dy = F.conv2d(F.pad(gray_t, (1, 1, 1, 1), mode='replicate'), prewitt_y)
|
|
|
|
elif method == "laplacian":
|
|
laplacian = torch.tensor([[[0, -1, 0], [-1, 4, -1], [0, -1, 0]]],
|
|
dtype=gray.dtype, device=gray.device).view(1, 1, 3, 3)
|
|
|
|
edges = F.conv2d(F.pad(gray_t, (1, 1, 1, 1), mode='replicate'), laplacian)
|
|
edges = torch.abs(edges)
|
|
|
|
else: # canny_approx
|
|
# Simplified Canny (Gaussian + Sobel)
|
|
# Gaussian blur first
|
|
sigma = 1.0
|
|
kernel_size = 5
|
|
x = torch.arange(kernel_size, dtype=gray.dtype, device=gray.device) - kernel_size // 2
|
|
kernel_1d = torch.exp(-x ** 2 / (2 * sigma ** 2))
|
|
kernel_1d = kernel_1d / kernel_1d.sum()
|
|
|
|
kernel_h = kernel_1d.view(1, 1, 1, -1)
|
|
blurred = F.conv2d(F.pad(gray_t, (2, 2, 0, 0), mode='replicate'), kernel_h)
|
|
|
|
kernel_v = kernel_1d.view(1, 1, -1, 1)
|
|
blurred = F.conv2d(F.pad(blurred, (0, 0, 2, 2), mode='replicate'), kernel_v)
|
|
|
|
# Sobel
|
|
sobel_x = torch.tensor([[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]],
|
|
dtype=gray.dtype, device=gray.device).view(1, 1, 3, 3)
|
|
sobel_y = torch.tensor([[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]],
|
|
dtype=gray.dtype, device=gray.device).view(1, 1, 3, 3)
|
|
|
|
dx = F.conv2d(F.pad(blurred, (1, 1, 1, 1), mode='replicate'), sobel_x)
|
|
dy = F.conv2d(F.pad(blurred, (1, 1, 1, 1), mode='replicate'), sobel_y)
|
|
|
|
# Calculate edge magnitude (except for laplacian)
|
|
if method != "laplacian":
|
|
edges = torch.sqrt(dx * dx + dy * dy)
|
|
|
|
edges = edges.permute(0, 2, 3, 1)
|
|
|
|
# Normalize
|
|
edges = (edges - edges.min()) / (edges.max() - edges.min() + 1e-7)
|
|
|
|
# Apply strength and threshold
|
|
edges = edges * strength
|
|
edges = torch.clamp((edges - threshold) / (1.0 - threshold + 1e-7), 0.0, 1.0)
|
|
|
|
# Invert if requested
|
|
if invert:
|
|
edges = 1.0 - edges
|
|
|
|
# Convert to 3-channel
|
|
edges = edges.repeat(1, 1, 1, 3)
|
|
|
|
print(f"✓ Edges detected")
|
|
print(f" Method: {method}")
|
|
print(f" Strength: {strength:.2f}")
|
|
print(f" Threshold: {threshold:.2f}")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (edges,)
|
|
|
|
|
|
class ImageEnhancement:
|
|
"""
|
|
General image enhancement
|
|
Sharpen, contrast, vibrance
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"sharpen": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.1,
|
|
"display": "number"
|
|
}),
|
|
"contrast": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.1,
|
|
"display": "number"
|
|
}),
|
|
"vibrance": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": -1.0,
|
|
"max": 1.0,
|
|
"step": 0.1,
|
|
"display": "number"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("enhanced",)
|
|
FUNCTION = "enhance"
|
|
CATEGORY = "Texture Alchemist/Filters"
|
|
|
|
def enhance(self, image, sharpen, contrast, vibrance):
|
|
"""Apply enhancements"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Image Enhancement")
|
|
print(f"{'='*60}")
|
|
|
|
result = image.clone()
|
|
|
|
# Sharpen
|
|
if sharpen > 0.0:
|
|
kernel = torch.tensor([[[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]]],
|
|
dtype=image.dtype, device=image.device).view(1, 1, 3, 3)
|
|
kernel = kernel / 9.0 # Normalize
|
|
|
|
image_t = result.permute(0, 3, 1, 2)
|
|
|
|
sharpened = torch.zeros_like(image_t)
|
|
for c in range(image_t.shape[1]):
|
|
sharpened[:, c:c+1, :, :] = F.conv2d(
|
|
F.pad(image_t[:, c:c+1, :, :], (1, 1, 1, 1), mode='replicate'),
|
|
kernel
|
|
)
|
|
|
|
sharpened = sharpened.permute(0, 2, 3, 1)
|
|
result = result + (sharpened - result) * sharpen
|
|
|
|
# Contrast
|
|
if contrast != 1.0:
|
|
mean = result.mean()
|
|
result = (result - mean) * contrast + mean
|
|
|
|
# Vibrance (smart saturation)
|
|
if vibrance != 0.0:
|
|
# Convert to HSV
|
|
r, g, b = result[:, :, :, 0], result[:, :, :, 1], result[:, :, :, 2]
|
|
|
|
max_c = torch.max(torch.max(r, g), b)
|
|
min_c = torch.min(torch.min(r, g), b)
|
|
|
|
# Saturation
|
|
s = torch.where(max_c != 0, (max_c - min_c) / (max_c + 1e-7), torch.zeros_like(max_c))
|
|
|
|
# Vibrance affects less saturated colors more
|
|
vibrance_mask = (1.0 - s) * torch.abs(vibrance)
|
|
|
|
# Adjust saturation
|
|
new_s = s + vibrance_mask * torch.sign(vibrance)
|
|
new_s = torch.clamp(new_s, 0.0, 1.0)
|
|
|
|
# Convert back (simplified)
|
|
scale = new_s / (s + 1e-7)
|
|
result = torch.stack([
|
|
min_c + (r - min_c) * scale,
|
|
min_c + (g - min_c) * scale,
|
|
min_c + (b - min_c) * scale
|
|
], dim=3)
|
|
|
|
result = torch.clamp(result, 0.0, 1.0)
|
|
|
|
print(f"✓ Enhancement applied")
|
|
print(f" Sharpen: {sharpen:.2f}")
|
|
print(f" Contrast: {contrast:.2f}")
|
|
print(f" Vibrance: {vibrance:.2f}")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (result,)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DenoiseFilter": DenoiseFilter,
|
|
"EdgeDetection": EdgeDetection,
|
|
"ImageEnhancement": ImageEnhancement,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DenoiseFilter": "Denoise Filter",
|
|
"EdgeDetection": "Edge Detection",
|
|
"ImageEnhancement": "Image Enhancement",
|
|
}
|
|
|