utils/color_fix.py

"""
Color correction utilities for super-resolution outputs.

Three methods are provided to align the color statistics of SR results
back to the input LR image:

- AdaIN: Global adaptive instance normalization (mean/std matching)
- Wavelet: Multi-level wavelet decomposition, replacing low-frequency color
- YCbCr: Direct color space replacement (Y from SR, CbCr from LR)

Original wavelet/AdaIN reference:
    https://github.com/pkuliyi2015/sd-webui-stablesr/blob/master/srmodule/colorfix.py
"""

import torch
from PIL import Image
from torch import Tensor
from torch.nn import functional as F
from torchvision.transforms import ToTensor, ToPILImage


# ---------------------------------------------------------------------------
# YCbCr color-space conversions (PyTorch, range [0,1])
# ---------------------------------------------------------------------------

def rgb2ycbcr_torch(im: Tensor, only_y: bool = True) -> Tensor:
    """MATLAB-compatible rgb2ycbcr.

    Args:
        im:  float tensor, N×3×H×W, values in [0, 1].
        only_y:  if True, return only the Y channel (N×1×H×W).

    Returns:
        YCbCr tensor in [0, 1].
    """
    im_temp = im.permute(0, 2, 3, 1) * 255.0  # N×H×W×3

    if only_y:
        coeff = torch.tensor([65.481, 128.553, 24.966],
                             device=im.device, dtype=im.dtype).view(3, 1) / 255.0
        result = torch.matmul(im_temp, coeff) + 16.0
    else:
        scale = torch.tensor(
            [[65.481, -37.797, 112.0],
             [128.553, -74.203, -93.786],
             [24.966, 112.0, -18.214]],
            device=im.device, dtype=im.dtype
        ) / 255.0
        bias = torch.tensor([16, 128, 128],
                            device=im.device, dtype=im.dtype).view(1, 1, 1, 3)
        result = torch.matmul(im_temp, scale) + bias

    result = result / 255.0
    result = result.clamp_(0.0, 1.0)
    return result.permute(0, 3, 1, 2)


def ycbcr2rgb_torch(im: Tensor) -> Tensor:
    """MATLAB-compatible ycbcr2rgb.

    Args:
        im:  float tensor, N×3×H×W, values in [0, 1].

    Returns:
        RGB tensor in [0, 1].
    """
    im_temp = im.permute(0, 2, 3, 1) * 255.0  # N×H×W×3

    scale = torch.tensor(
        [[0.00456621, 0.00456621, 0.00456621],
         [0, -0.00153632, 0.00791071],
         [0.00625893, -0.00318811, 0]],
        device=im.device, dtype=im.dtype
    ) * 255.0
    bias = torch.tensor(
        [-222.921, 135.576, -276.836],
        device=im.device, dtype=im.dtype
    ).view(1, 1, 1, 3)

    result = torch.matmul(im_temp, scale) + bias
    result = result / 255.0
    result = result.clamp_(0.0, 1.0)
    return result.permute(0, 3, 1, 2)


# ---------------------------------------------------------------------------
# AdaIN color fix
# ---------------------------------------------------------------------------

def _calc_mean_std(feat: Tensor, eps: float = 1e-5):
    """Compute per-channel mean and std for a 4D tensor."""
    b, c = feat.size()[:2]
    feat_var = feat.reshape(b, c, -1).var(dim=2) + eps
    feat_std = feat_var.sqrt().reshape(b, c, 1, 1)
    feat_mean = feat.reshape(b, c, -1).mean(dim=2).reshape(b, c, 1, 1)
    return feat_mean, feat_std


def _adain(content: Tensor, style: Tensor) -> Tensor:
    """Adaptive instance normalization: match content stats to style stats."""
    style_mean, style_std = _calc_mean_std(style)
    content_mean, content_std = _calc_mean_std(content)
    size = content.size()
    normalized = (content - content_mean.expand(size)) / content_std.expand(size)
    return normalized * style_std.expand(size) + style_mean.expand(size)


def adain_color_fix(target: Image.Image, source: Image.Image) -> Image.Image:
    """Globally match the mean & std of *target* to *source* via AdaIN.

    Args:
        target:  SR output (PIL Image).
        source:  LR input (PIL Image).

    Returns:
        Color-corrected PIL Image.
    """
    target_t = ToTensor()(target).unsqueeze(0)
    source_t = ToTensor()(source).unsqueeze(0)
    result_t = _adain(target_t, source_t)
    result_t = result_t.squeeze(0).clamp_(0.0, 1.0)
    return ToPILImage()(result_t)


# ---------------------------------------------------------------------------
# Wavelet color fix
# ---------------------------------------------------------------------------

_WAVELET_KERNEL = torch.tensor(
    [[0.0625, 0.125, 0.0625],
     [0.125, 0.25, 0.125],
     [0.0625, 0.125, 0.0625]]
)


def _wavelet_blur(image: Tensor, radius: int) -> Tensor:
    """Low-pass filter via dilated Gaussian-like conv.

    Args:
        image:  (1, 3, H, W) tensor.
        radius: dilation radius.

    Returns:
        Blurred tensor of same shape.
    """
    kernel = _WAVELET_KERNEL.to(dtype=image.dtype, device=image.device)
    kernel = kernel[None, None].repeat(3, 1, 1, 1)  # (3,1,3,3)
    padded = F.pad(image, (radius, radius, radius, radius), mode='replicate')
    return F.conv2d(padded, kernel, groups=3, dilation=radius)


def _wavelet_decomp(image: Tensor, levels: int = 5):
    """Multi-level wavelet decomposition.

    Returns:
        (high_freq, low_freq): both (1, 3, H, W).
    """
    high_freq = torch.zeros_like(image)
    cur = image
    for i in range(levels):
        radius = 2 ** i
        low = _wavelet_blur(cur, radius)
        high_freq = high_freq + (cur - low)
        cur = low
    return high_freq, cur  # high, low


def _wavelet_reconstruct(content: Tensor, style: Tensor) -> Tensor:
    """Replace low-frequency of *content* with that of *style*."""
    content_high, _ = _wavelet_decomp(content)
    _, style_low = _wavelet_decomp(style)
    return content_high + style_low


def wavelet_color_fix(target: Image.Image, source: Image.Image) -> Image.Image:
    """Multi-level wavelet color correction.

    Replaces the low-frequency component of *target* (SR) with that of
    *source* (LR), preserving SR high-frequency details.

    Args:
        target:  SR output (PIL Image).
        source:  LR input (PIL Image).

    Returns:
        Color-corrected PIL Image.
    """
    target_t = ToTensor()(target).unsqueeze(0)
    source_t = ToTensor()(source).unsqueeze(0)
    result_t = _wavelet_reconstruct(target_t, source_t)
    result_t = result_t.squeeze(0).clamp_(0.0, 1.0)
    return ToPILImage()(result_t)


# ---------------------------------------------------------------------------
# YCbCr color fix
# ---------------------------------------------------------------------------

def ycbcr_color_fix(target: Image.Image, source: Image.Image) -> Image.Image:
    """Replace Cb/Cr channels of *target* with those of *source*.

    Keeps SR luminance (Y) and borrows LR chrominance (CbCr), which is
    the most direct way to restore color fidelity.

    Args:
        target:  SR output (PIL Image).
        source:  LR input (PIL Image).

    Returns:
        Color-corrected PIL Image.
    """
    target_t = ToTensor()(target).unsqueeze(0)
    source_t = ToTensor()(source).unsqueeze(0)

    content_y = rgb2ycbcr_torch(target_t, only_y=True)          # (1,1,H,W)
    style_ycbcr = rgb2ycbcr_torch(source_t, only_y=False)      # (1,3,H,W)

    combined = torch.cat([content_y, style_ycbcr[:, 1:]], dim=1)  # Y_sr + CbCr_lr
    result_t = ycbcr2rgb_torch(combined)

    result_t = result_t.squeeze(0).clamp_(0.0, 1.0)
    return ToPILImage()(result_t)


# ---------------------------------------------------------------------------
# Unified entry
# ---------------------------------------------------------------------------

COLOR_FIX_METHODS = {
    'adain': adain_color_fix,
    'wavelet': wavelet_color_fix,
    'ycbcr': ycbcr_color_fix,
    'none': lambda target, source: target,
}


def apply_color_fix(target: Image.Image, source: Image.Image,
                    method: str = 'ycbcr') -> Image.Image:
    """Apply the specified color correction method.

    Args:
        target:  SR output image.
        source:  LR input image.
        method:  one of 'adain', 'wavelet', 'ycbcr', 'none'.

    Returns:
        Color-corrected image.
    """
    if method not in COLOR_FIX_METHODS:
        raise ValueError(
            f"Unknown color fix method: '{method}'. "
            f"Available: {list(COLOR_FIX_METHODS.keys())}"
        )
    return COLOR_FIX_METHODS[method](target, source)