# IMPORT LIBS # Import pyTorch import torch # Import Tensor type for type hinting from torch import Tensor # Python Imaging Library for image operations from PIL import Image, ImageOps # PyTorch's functional API for neural network operations from torch.nn import functional as F # Transformations for image processing from torchvision.transforms.v2 import ToTensor, ToPILImage # From util import numpy as np #import PIL.Image as Image # ========= pil2tensor ========= # def pil2tensor(image): print("[DICKSON-NODES] pil2tensor") return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) # ========= tensor2pil ========= # def tensor2pil(image): print("[DICKSON-NODES] tensor2pil") return Image.fromarray( np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) ) # Main function for AdaIN color matching def adain_color_match(target: Image, source: Image): print("[DICKSON-NODES] adain_color_match") # Convert images to tensors to_tensor = ToTensor() print("[DICKSON-NODES] ToTensor") target_tensor = to_tensor(target).unsqueeze(0) print("[DICKSON-NODES] target ToTensor") source_tensor = to_tensor(source).unsqueeze(0) print("[DICKSON-NODES] source ToTensor") # Apply adaptive instance normalization result_tensor = adaptive_instance_normalization(target_tensor, source_tensor) # Convert tensor back to image print("[DICKSON-NODES] ToPILImage") to_image = ToPILImage() result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0)) return result_image # Main function for wavelet color matching def wavelet_color_match(target: Image, source: Image): # Print a message to indicate the function has been called print("[DICKSON-NODES] wavelet_color_match") # Check if CUDA is available #device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Check for available GPU devices if torch.cuda.is_available(): # If CUDA (NVIDIA GPU) is available, use it device = torch.device("cuda") elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): # If MPS (Apple Silicon GPU) is available, use it device = torch.device("mps") else: # If no GPU is available, use CPU device = torch.device("cpu") # Print which device (CPU or GPU) is being used print(f"[DICKSON-NODES] Using device: {device}") # Resize the source image to match the target image size source = source.resize(target.size, resample=Image.Resampling.LANCZOS) # Create a function to convert PIL Images to PyTorch tensors to_tensor = ToTensor() #target_tensor = to_tensor(target).unsqueeze(0) #source_tensor = to_tensor(source).unsqueeze(0) # Convert target image to tensor, add a batch dimension, and move to the selected device target_tensor = to_tensor(target).unsqueeze(0).to(device) # Convert source image to tensor, add a batch dimension, and move to the selected device source_tensor = to_tensor(source).unsqueeze(0).to(device) # Apply wavelet reconstruction result_tensor = wavelet_reconstruction(target_tensor, source_tensor) # Move the result tensor back to CPU for image conversion result_tensor = result_tensor.cpu() # Create a function to convert PyTorch tensors back to PIL Images to_image = ToPILImage() # Convert the result tensor to a PIL Image, removing the batch dimension and clamping values between 0 and 1 result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0)) return result_image # Function to calculate mean and standard deviation of a tensor def calc_mean_std(feat: Tensor, eps=1e-5): """ Calculate mean and std for adaptive_instance_normalization. Args: feat (Tensor): 4D tensor. eps (float): A small value added to the variance to avoid divide-by-zero. Default: 1e-5. """ print("[DICKSON-NODES] calc_mean_std") size = feat.size() # Get the size of the input tensor assert len(size) == 4, 'The input feature should be 4D tensor.' # Ensure input is 4D b, c = size[:2] # Get batch size and number of channels feat_var = feat.view(b, c, -1).var(dim=2) + eps # Calculate variance feat_std = feat_var.sqrt().view(b, c, 1, 1) # Calculate standard deviation feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1) # Calculate mean return feat_mean, feat_std # Return mean and standard deviation # Function for adaptive instance normalization def adaptive_instance_normalization(content_feat:Tensor, style_feat:Tensor): """ Adaptive instance normalization. Adjust the reference features to have the similar color and illuminations as those in the degradate features. Args: content_feat (Tensor): The reference feature. style_feat (Tensor): The degradate features. """ print("[DICKSON-NODES] adaptive_instance_normalization") size = content_feat.size() # Get size of content feature tensor style_mean, style_std = calc_mean_std(style_feat) # Calculate mean and std of style content_mean, content_std = calc_mean_std(content_feat) # Calculate mean and std of content normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size) # Normalize content feature return normalized_feat * style_std.expand(size) + style_mean.expand(size) # Apply style's mean and std to normalized content # Function for wavelet blur def wavelet_blur(image: Tensor, radius: int): """ Apply wavelet blur to the input tensor. """ # Print a message to indicate the function has been called print("[DICKSON-NODES] wavelet_blur") # Get the device (CPU or GPU) that the input image is on # Ensure the image is on the correct device device = image.device print(f"[DICKSON-NODES] Using device: {device}") # Define the blur kernel values # input shape: (1, 3, H, W) # convolution kernel kernel_vals = [ [0.0625, 0.125, 0.0625], [0.125, 0.25, 0.125], [0.0625, 0.125, 0.0625], ] # Convert the kernel values to a PyTorch tensor on the same device as the input image kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device) # Add two dimensions to the kernel to make it a 4D tensor (required for conv2d) # add channel dimensions to the kernel to make it a 4D tensor kernel = kernel[None, None] # Repeat the kernel for each input channel (assuming 3 channels: R, G, B) # repeat the kernel across all input channels kernel = kernel.repeat(3, 1, 1, 1) # Pad the input image to maintain size after convolution image = F.pad(image, (radius, radius, radius, radius), mode='replicate') # Apply convolution to blur the image output = F.conv2d(image, kernel, groups=3, dilation=radius) # Return the blurred image return output # Function for wavelet decomposition def wavelet_decomposition(image: Tensor, levels=5): """ Apply wavelet decomposition to the input tensor. This function only returns the low frequency & the high frequency. """ print("[DICKSON-NODES] wavelet_decomposition") # Initialize high_freq tensor with zeros, same shape as input image high_freq = torch.zeros_like(image) # Iterate through each level of decomposition for i in range(levels): # Calculate the blur radius for this level radius = 2 ** i # Apply wavelet blur to get low frequency component low_freq = wavelet_blur(image, radius) # Calculate and accumulate high frequency component high_freq += (image - low_freq) # Update image for next iteration image = low_freq # Return high frequency and low frequency components return high_freq, low_freq # Function for wavelet reconstruction def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor): """ Apply wavelet decomposition, so that the content will have the same color as the style. """ print("[DICKSON-NODES] wavelet_reconstruction") # calculate the wavelet decomposition of the content feature content_high_freq, content_low_freq = wavelet_decomposition(content_feat) del content_low_freq # calculate the wavelet decomposition of the style feature style_high_freq, style_low_freq = wavelet_decomposition(style_feat) del style_high_freq # reconstruct the content feature with the style's high frequency return content_high_freq + style_low_freq