Dynamatrix/DiffBIR-OpenXLab

1'''2# --------------------------------------------------------------------------------3#   Color fixed script from Li Yi (https://github.com/pkuliyi2015/sd-webui-stablesr/blob/master/srmodule/colorfix.py)4# --------------------------------------------------------------------------------5'''6 7import torch8from PIL import Image9from torch import Tensor10from torch.nn import functional as F11from torchvision.transforms import ToTensor, ToPILImage12 13 14def adain_color_fix(target: Image, source: Image):15    # Convert images to tensors16    to_tensor = ToTensor()17    target_tensor = to_tensor(target).unsqueeze(0)18    source_tensor = to_tensor(source).unsqueeze(0)19 20    # Apply adaptive instance normalization21    result_tensor = adaptive_instance_normalization(target_tensor, source_tensor)22 23    # Convert tensor back to image24    to_image = ToPILImage()25    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))26 27    return result_image28 29def wavelet_color_fix(target: Image, source: Image):30    # Convert images to tensors31    to_tensor = ToTensor()32    target_tensor = to_tensor(target).unsqueeze(0)33    source_tensor = to_tensor(source).unsqueeze(0)34 35    # Apply wavelet reconstruction36    result_tensor = wavelet_reconstruction(target_tensor, source_tensor)37 38    # Convert tensor back to image39    to_image = ToPILImage()40    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))41 42    return result_image43 44def calc_mean_std(feat: Tensor, eps=1e-5):45    """Calculate mean and std for adaptive_instance_normalization.46    Args:47        feat (Tensor): 4D tensor.48        eps (float): A small value added to the variance to avoid49            divide-by-zero. Default: 1e-5.50    """51    size = feat.size()52    assert len(size) == 4, 'The input feature should be 4D tensor.'53    b, c = size[:2]54    feat_var = feat.reshape(b, c, -1).var(dim=2) + eps55    feat_std = feat_var.sqrt().reshape(b, c, 1, 1)56    feat_mean = feat.reshape(b, c, -1).mean(dim=2).reshape(b, c, 1, 1)57    return feat_mean, feat_std58 59def adaptive_instance_normalization(content_feat:Tensor, style_feat:Tensor):60    """Adaptive instance normalization.61    Adjust the reference features to have the similar color and illuminations62    as those in the degradate features.63    Args:64        content_feat (Tensor): The reference feature.65        style_feat (Tensor): The degradate features.66    """67    size = content_feat.size()68    style_mean, style_std = calc_mean_std(style_feat)69    content_mean, content_std = calc_mean_std(content_feat)70    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)71    return normalized_feat * style_std.expand(size) + style_mean.expand(size)72 73def wavelet_blur(image: Tensor, radius: int):74    """75    Apply wavelet blur to the input tensor.76    """77    # input shape: (1, 3, H, W)78    # convolution kernel79    kernel_vals = [80        [0.0625, 0.125, 0.0625],81        [0.125, 0.25, 0.125],82        [0.0625, 0.125, 0.0625],83    ]84    kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)85    # add channel dimensions to the kernel to make it a 4D tensor86    kernel = kernel[None, None]87    # repeat the kernel across all input channels88    kernel = kernel.repeat(3, 1, 1, 1)89    image = F.pad(image, (radius, radius, radius, radius), mode='replicate')90    # apply convolution91    output = F.conv2d(image, kernel, groups=3, dilation=radius)92    return output93 94def wavelet_decomposition(image: Tensor, levels=5):95    """96    Apply wavelet decomposition to the input tensor.97    This function only returns the low frequency & the high frequency.98    """99    high_freq = torch.zeros_like(image)100    for i in range(levels):101        radius = 2 ** i102        low_freq = wavelet_blur(image, radius)103        high_freq += (image - low_freq)104        image = low_freq105 106    return high_freq, low_freq107 108def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor):109    """110    Apply wavelet decomposition, so that the content will have the same color as the style.111    """112    # calculate the wavelet decomposition of the content feature113    content_high_freq, content_low_freq = wavelet_decomposition(content_feat)114    del content_low_freq115    # calculate the wavelet decomposition of the style feature116    style_high_freq, style_low_freq = wavelet_decomposition(style_feat)117    del style_high_freq118    # reconstruct the content feature with the style's high frequency119    return content_high_freq + style_low_freq