diff --git a/__init__.py b/__init__.py index ef08fa8..6b62925 100644 --- a/__init__.py +++ b/__init__.py @@ -8,7 +8,7 @@ from .stream_cfg import StreamCFG from .stream_conditioning import StreamConditioning from .particle_nodes import DepthMapWarpNode import re - +from .image_enhance import UnsharpMask NODE_CLASS_MAPPINGS = { @@ -39,6 +39,7 @@ NODE_CLASS_MAPPINGS = { "FastWebcamCapture": FastWebcamCapture, "MultilineText": MultilineText, "LoadImageFromPath_": LoadImageFromPath_, + "UnsharpMask": UnsharpMask, } NODE_DISPLAY_NAME_MAPPINGS = {} diff --git a/image_enhance.py b/image_enhance.py new file mode 100644 index 0000000..1692eca --- /dev/null +++ b/image_enhance.py @@ -0,0 +1,496 @@ +import torch +import torch.nn.functional as F + +class HistogramStretch: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "clip_percent": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 0.5, "step": 0.05}), + "amount": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "enhance" + CATEGORY = "image/enhance" + + def enhance(self, image, clip_percent, amount): + # Automatically stretch histogram for better contrast + B, H, W, C = image.shape + x = image.permute(0, 3, 1, 2) # to BCHW + + # Calculate histogram bounds + low = torch.quantile(x, clip_percent, dim=2, keepdim=True) + low = torch.quantile(low, clip_percent, dim=3, keepdim=True) + + high = torch.quantile(x, 1.0 - clip_percent, dim=2, keepdim=True) + high = torch.quantile(high, 1.0 - clip_percent, dim=3, keepdim=True) + + # Stretch contrast with amount control + enhanced = x + amount * ((x - low) / (high - low) - x) + enhanced = torch.clamp(enhanced, 0, 1) + + return (enhanced.permute(0, 2, 3, 1),) + +class AdaptiveSharpen: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "max_amount": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}), + "edge_threshold": ("FLOAT", {"default": 5.0, "min": 1.0, "max": 10.0, "step": 0.5}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "enhance" + CATEGORY = "image/enhance" + + def enhance(self, image, max_amount, edge_threshold): + # Content-aware sharpening using Laplacian + B, H, W, C = image.shape + x = image.permute(0, 3, 1, 2) # to BCHW + + # Laplacian kernel for edge detection + laplacian = torch.tensor([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=torch.float32, device=x.device).view(1, 1, 3, 3) + + # Detect edges + edges = F.conv2d(x, laplacian.expand(C, 1, 3, 3), padding=1, groups=C) + edges = torch.abs(edges) + + # Adaptive amount based on edge strength + amount = torch.sigmoid(edges * edge_threshold) * max_amount + + # Apply sharpening + enhanced = x + amount * edges + enhanced = torch.clamp(enhanced, 0, 1) + + return (enhanced.permute(0, 2, 3, 1),) + +class ChannelBalance: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "amount": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "enhance" + CATEGORY = "image/enhance" + + def enhance(self, image, amount): + # Automatic RGB channel balancing + B, H, W, C = image.shape + + # Calculate mean for each channel + means = image.mean(dim=(0, 1, 2), keepdim=True) + target_mean = means.mean() + + # Adjust each channel to match target mean with amount control + balance = (target_mean / means) + enhanced = image + amount * (image * balance - image) + enhanced = torch.clamp(enhanced, 0, 1) + + return (enhanced,) + +class MultiScaleDetail: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "small_amount": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.005}), + "medium_amount": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.005}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "enhance" + CATEGORY = "image/enhance" + + def enhance(self, image, small_amount, medium_amount): + # Multi-scale detail enhancement with zero-sum kernels + B, H, W, C = image.shape + x = image.permute(0, 3, 1, 2) # to BCHW + + # Small details (3x3) - zero-sum kernel + small_kernel = torch.tensor([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype=torch.float32, device=x.device).view(1, 1, 3, 3) + small_details = F.conv2d(x, small_kernel.expand(C, 1, 3, 3), padding=1, groups=C) + + # Medium details (5x5) - zero-sum kernel + medium_kernel = torch.tensor([[-1, -1, -1, -1, -1], [-1, 2, 2, 2, -1], [-1, 2, 8, 2, -1], [-1, 2, 2, 2, -1], [-1, -1, -1, -1, -1]], dtype=torch.float32, device=x.device).view(1, 1, 5, 5) + medium_kernel = medium_kernel - medium_kernel.mean() # Ensure zero-sum + medium_details = F.conv2d(x, medium_kernel.expand(C, 1, 5, 5), padding=2, groups=C) + + # Combine details with adjustable weighting + enhanced = x + small_amount * small_details + medium_amount * medium_details + enhanced = torch.clamp(enhanced, 0, 1) + + return (enhanced.permute(0, 2, 3, 1),) + +class UnsharpMask: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "blur_radius": ("FLOAT", {"default": 1.2, "min": 0.1, "max": 10.0, "step": 0.1}), + "amount": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 3.0, "step": 0.05}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "sharpen" + CATEGORY = "image/enhance" + + def sharpen(self, image, blur_radius, amount): + # Input: (B, H, W, C) + B, H, W, C = image.shape + x = image.permute(0, 3, 1, 2) # to BCHW + + # Gaussian kernel (separable) + def gaussian_kernel1d(radius, sigma): + size = int(2 * radius + 1) + coords = torch.arange(size) - radius + kernel = torch.exp(-coords**2 / (2 * sigma**2)) + kernel /= kernel.sum() + return kernel.to(x.device) + + radius = int(blur_radius) + kernel = gaussian_kernel1d(radius, sigma=blur_radius) + + # Create vertical and horizontal kernels + kernel_v = kernel.view(1, 1, -1, 1) + kernel_h = kernel.view(1, 1, 1, -1) + + # Blur H and W + blurred = F.conv2d(x, kernel_v.expand(C, 1, -1, 1), padding=(0, radius), groups=C) + blurred = F.conv2d(blurred, kernel_h.expand(C, 1, 1, -1), padding=(radius, 0), groups=C) + + # Unsharp mask + sharpened = x + amount * (x - blurred) + sharpened = torch.clamp(sharpened, 0, 1) + + return (sharpened.permute(0, 2, 3, 1),) + +class ICAS: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "amount": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0}), + "radius": ("INT", {"default": 2, "min": 1, "max": 5}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "enhance" + CATEGORY = "image/enhance" + + def enhance(self, image, amount, radius): + # Image Contrast Adaptive Sharpen + B, H, W, C = image.shape + x = image.permute(0, 3, 1, 2) # to BCHW + + # Efficient Gaussian blur using separable kernels + def gaussian_kernel1d(radius, sigma): + size = int(2 * radius + 1) + coords = torch.arange(size) - radius + kernel = torch.exp(-coords**2 / (2 * sigma**2)) + kernel /= kernel.sum() + return kernel.to(x.device) + + # Create and apply Gaussian blur + kernel = gaussian_kernel1d(radius, sigma=radius) + kernel_v = kernel.view(1, 1, -1, 1) + kernel_h = kernel.view(1, 1, 1, -1) + + # Apply separable blur + blurred = F.conv2d(x, kernel_v.expand(C, 1, -1, 1), padding=(0, radius), groups=C) + blurred = F.conv2d(blurred, kernel_h.expand(C, 1, 1, -1), padding=(radius, 0), groups=C) + + # Calculate local contrast + local_contrast = torch.abs(x - blurred) + local_contrast = F.avg_pool2d(local_contrast, kernel_size=3, stride=1, padding=1) + + # Adaptive sharpening mask + mask = torch.sigmoid(local_contrast * 10.0) # Adjust sensitivity + mask = mask * amount + + # Apply sharpening + sharpened = x + mask * (x - blurred) + sharpened = torch.clamp(sharpened, 0, 1) + + return (sharpened.permute(0, 2, 3, 1),) + +class FastBilateralDetail: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "amount": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0}), + "radius": ("INT", {"default": 3, "min": 1, "max": 7}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "enhance" + CATEGORY = "image/enhance" + + def enhance(self, image, amount, radius): + # Fast bilateral filter-based detail enhancement + B, H, W, C = image.shape + x = image.permute(0, 3, 1, 2) # to BCHW + + # Efficient bilateral filter approximation using separable Gaussian + def gaussian_kernel1d(radius, sigma): + size = int(2 * radius + 1) + coords = torch.arange(size) - radius + kernel = torch.exp(-coords**2 / (2 * sigma**2)) + kernel /= kernel.sum() + return kernel.to(x.device) + + # Create kernels + kernel = gaussian_kernel1d(radius, sigma=radius) + kernel_v = kernel.view(1, 1, -1, 1) + kernel_h = kernel.view(1, 1, 1, -1) + + # Apply separable Gaussian blur + blurred = F.conv2d(x, kernel_v.expand(C, 1, -1, 1), padding=(0, radius), groups=C) + blurred = F.conv2d(blurred, kernel_h.expand(C, 1, 1, -1), padding=(radius, 0), groups=C) + + # Calculate edge strength using Sobel + sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32, device=x.device).view(1, 1, 3, 3) + sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32, device=x.device).view(1, 1, 3, 3) + + grad_x = F.conv2d(x, sobel_x.expand(C, 1, 3, 3), padding=1, groups=C) + grad_y = F.conv2d(x, sobel_y.expand(C, 1, 3, 3), padding=1, groups=C) + edge_strength = torch.sqrt(grad_x**2 + grad_y**2) + edge_strength = F.avg_pool2d(edge_strength, kernel_size=3, stride=1, padding=1) + + # Create adaptive mask based on edge strength + mask = torch.sigmoid(edge_strength * 5.0) # Adjust sensitivity + mask = mask * amount + + # Apply detail enhancement + enhanced = x + mask * (x - blurred) + enhanced = torch.clamp(enhanced, 0, 1) + + return (enhanced.permute(0, 2, 3, 1),) + +class WaveletDetail: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "amount": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0}), + "scale": ("INT", {"default": 2, "min": 1, "max": 4}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "enhance" + CATEGORY = "image/enhance" + + def enhance(self, image, amount, scale): + # Wavelet-based detail enhancement + B, H, W, C = image.shape + x = image.permute(0, 3, 1, 2) # to BCHW + + # Haar wavelet decomposition + def haar_decompose(x): + # Downsample + x_low = F.avg_pool2d(x, kernel_size=2, stride=2) + # Upsample and subtract for detail + x_high = x - F.interpolate(x_low, size=x.shape[2:], mode='bilinear', align_corners=False) + return x_low, x_high + + # Multi-scale decomposition + details = [] + current = x + for _ in range(scale): + current, detail = haar_decompose(current) + # Upsample detail to match original size + detail = F.interpolate(detail, size=(H, W), mode='bilinear', align_corners=False) + details.append(detail) + + # Enhance details + enhanced = x + for detail in details: + enhanced = enhanced + amount * detail + + enhanced = torch.clamp(enhanced, 0, 1) + return (enhanced.permute(0, 2, 3, 1),) + +class FrequencySelective: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "amount": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0}), + "radius": ("INT", {"default": 3, "min": 1, "max": 7}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "enhance" + CATEGORY = "image/enhance" + + def enhance(self, image, amount, radius): + # Frequency-selective detail enhancement + B, H, W, C = image.shape + x = image.permute(0, 3, 1, 2) # to BCHW + + # Efficient Gaussian blur using separable kernels + def gaussian_kernel1d(radius, sigma): + size = int(2 * radius + 1) + coords = torch.arange(size) - radius + kernel = torch.exp(-coords**2 / (2 * sigma**2)) + kernel /= kernel.sum() + return kernel.to(x.device) + + # Create kernels + kernel = gaussian_kernel1d(radius, sigma=radius) + kernel_v = kernel.view(1, 1, -1, 1) + kernel_h = kernel.view(1, 1, 1, -1) + + # Apply separable Gaussian blur + blurred = F.conv2d(x, kernel_v.expand(C, 1, -1, 1), padding=(0, radius), groups=C) + blurred = F.conv2d(blurred, kernel_h.expand(C, 1, 1, -1), padding=(radius, 0), groups=C) + + # Frequency separation + high_freq = x - blurred + low_freq = blurred + + # Enhance high frequencies + enhanced = low_freq + amount * high_freq + enhanced = torch.clamp(enhanced, 0, 1) + + return (enhanced.permute(0, 2, 3, 1),) + +class GuidedDetail: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "amount": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0}), + "radius": ("INT", {"default": 3, "min": 1, "max": 7}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "enhance" + CATEGORY = "image/enhance" + + def enhance(self, image, amount, radius): + # Guided filter-based detail enhancement + B, H, W, C = image.shape + x = image.permute(0, 3, 1, 2) # to BCHW + + # Efficient Gaussian blur using separable kernels + def gaussian_kernel1d(radius, sigma): + size = int(2 * radius + 1) + coords = torch.arange(size) - radius + kernel = torch.exp(-coords**2 / (2 * sigma**2)) + kernel /= kernel.sum() + return kernel.to(x.device) + + # Create kernels + kernel = gaussian_kernel1d(radius, sigma=radius) + kernel_v = kernel.view(1, 1, -1, 1) + kernel_h = kernel.view(1, 1, 1, -1) + + # Apply separable Gaussian blur + blurred = F.conv2d(x, kernel_v.expand(C, 1, -1, 1), padding=(0, radius), groups=C) + blurred = F.conv2d(blurred, kernel_h.expand(C, 1, 1, -1), padding=(radius, 0), groups=C) + + # Calculate local variance + local_mean = F.avg_pool2d(x, kernel_size=2*radius+1, stride=1, padding=radius) + local_var = F.avg_pool2d(x**2, kernel_size=2*radius+1, stride=1, padding=radius) - local_mean**2 + + # Create adaptive mask based on local variance + mask = torch.sigmoid(local_var * 10.0) # Adjust sensitivity + mask = mask * amount + + # Apply detail enhancement + enhanced = x + mask * (x - blurred) + enhanced = torch.clamp(enhanced, 0, 1) + + return (enhanced.permute(0, 2, 3, 1),) + +class ImageConcat: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "direction": (["horizontal", "vertical"],), + "spacing": ("INT", {"default": 0, "min": 0, "max": 100}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "concat" + CATEGORY = "image/enhance" + + def concat(self, image1, image2, direction, spacing): + # Ensure both images have the same number of channels + if image1.shape[-1] != image2.shape[-1]: + raise ValueError("Images must have the same number of channels") + + # Add spacing + if spacing > 0: + if direction == "horizontal": + spacer = torch.zeros((image1.shape[0], image1.shape[1], spacing, image1.shape[3]), device=image1.device) + return (torch.cat([image1, spacer, image2], dim=2),) + else: + spacer = torch.zeros((image1.shape[0], spacing, image1.shape[2], image1.shape[3]), device=image1.device) + return (torch.cat([image1, spacer, image2], dim=1),) + + # Concatenate without spacing + if direction == "horizontal": + return (torch.cat([image1, image2], dim=2),) + else: + return (torch.cat([image1, image2], dim=1),) + +NODE_CLASS_MAPPINGS = { + "HistogramStretch": HistogramStretch, + "AdaptiveSharpen": AdaptiveSharpen, + "ChannelBalance": ChannelBalance, + "MultiScaleDetail": MultiScaleDetail, + "UnsharpMask": UnsharpMask, + "ICAS": ICAS, + "FastBilateralDetail": FastBilateralDetail, + "WaveletDetail": WaveletDetail, + "FrequencySelective": FrequencySelective, + "GuidedDetail": GuidedDetail, + "ImageConcat": ImageConcat, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "HistogramStretch": "Histogram Stretch", + "AdaptiveSharpen": "Adaptive Sharpen", + "ChannelBalance": "Channel Balance", + "MultiScaleDetail": "Multi-Scale Detail", + "UnsharpMask": "Unsharp Mask", + "ICAS": "Image Contrast Adaptive Sharpen", + "FastBilateralDetail": "Fast Bilateral Detail", + "WaveletDetail": "Wavelet Detail", + "FrequencySelective": "Frequency Selective", + "GuidedDetail": "Guided Detail", + "ImageConcat": "Image Concat", +} + +