<filename>mmedit/datasets/pipelines/bsrgan_degradation.py
import warnings
warnings.filterwarnings('ignore')

import os
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import cv2
from scipy.interpolate import interp2d

from ..registry import PIPELINES


# blur
class BlurKernel2D(object):
    def check_kernel(self):
        if self.ksize[0] % 2 == 0 or self.ksize[1] % 2 == 0:
            raise ValueError(
                f'Expect kernel size to be odd (1, 3, 5, etc.), but got {self.ksize}')

    def check_input(self, x):
        return isinstance(x, np.ndarray) and x.ndim == 3

    def __call__(self, im):
        if self.check_input(im):
            return self.apply(im)
        elif im is None:
            raise ValueError(f'Encounter None data in {type(self)}')
        else:
            raise ValueError(f'Expect input image to be 3D-array (HWC), but got {im.ndim}D-array. {type(self)}')

    def apply(self, im):
        raise NotImplementedError


class IsotropicGaussianBlur2D(BlurKernel2D):
    def __init__(self, kernel_size, std):
        assert kernel_size % 2 == 1
        self.kernel_size = (kernel_size, kernel_size)
        self.std = std

    def apply(self, im):
        if self.check_input(im):
            return cv2.GaussianBlur(im, self.kernel_size, self.std, self.std,
                                    borderType=cv2.BORDER_REFLECT_101)
        else:
            raise ValueError


class AnisotropicGaussianBlur2D(BlurKernel2D):
    def __init__(self, kernel_size, var, angle, scale=1):
        assert kernel_size % 2 == 1
        k_size = np.array([kernel_size, kernel_size])

        # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
        lambda_1, lambda_2 = var

        # Set COV matrix using Lambdas and Theta
        LAMBDA = np.diag([lambda_1, lambda_2])
        Q = np.array([[np.cos(angle), -np.sin(angle)],
                      [np.sin(angle), np.cos(angle)]])
        SIGMA = Q @ LAMBDA @ Q.T
        INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]

        # shifting kernel for aligned image
        MU = k_size // 2 - 0.5 * (scale - 1)
        MU = MU[None, None, :, None]

        # Create meshgrid for Gaussian
        [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
        Z = np.stack([X, Y], 2)[:, :, :, None]

        # Calcualte Gaussian for every pixel of the kernel
        ZZ = Z - MU
        ZZ_t = ZZ.transpose(0, 1, 3, 2)
        raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ))

        # Normalize the kernel and return
        self.kernel = raw_kernel / np.sum(raw_kernel)

    def apply(self, im):
        return cv2.filter2D(im, -1, self.kernel)


# camera_sensor_noise
# =================================================================================
# unprocess
def prepare_ccm():
    matrix = np.array([
        [21492, -5753, -3253, 2306, 8110, -425, 7066, 1327, 2068],
        [9804, -2699, -1302, -5813, 13164, 3058, -2428, 3065, 8675],
        [6847, -614, -1114, -4669, 12737, 2139, -1197, 2487, 6846],
        [13510, -6199, -1244, -4430, 12736, 1865, -332, 1443, 5024],
        [11511, -4358, -1065, -6524, 13767, 3068, -1467, 1984, 6045],
        [7175, -1983, -658, -8075, 15555, 2718, -2171, 2502, 7457],
        [7713, -2059, -653, -3882, 11495, 2726, -710, 1332, 5958]
    ]) / 10000

    num_ccms = len(matrix)
    xyz2cams = np.array(matrix).reshape(-1, 3, 3)

    weights = np.random.uniform(1e-8, 1e8, (num_ccms, 1, 1))
    weights_sum = np.sum(weights, axis=0)
    xyz2cam = np.sum(xyz2cams * weights, axis=0) / weights_sum

    # Multiplies with RGB -> XYZ to get RGB -> Camera CCM.
    rgb2xyz = np.array([[0.4124564, 0.3575761, 0.1804375],
                        [0.2126729, 0.7151522, 0.0721750],
                        [0.0193339, 0.1191920, 0.9503041]])
    rgb2cam = np.matmul(xyz2cam, rgb2xyz)

    # Normalizes each row.
    rgb2cam = rgb2cam / np.sum(rgb2cam, axis=-1, keepdims=True)
    return rgb2cam


def random_gains():
    """Generates random gains for brightening and white balance."""
    # RGB gain represents brightening.
    rgb_gain = np.random.uniform(2 ** (-0.1), 2 ** (0.3))

    # Red and blue gains for white balance.
    red_gain = np.random.uniform(1.2, 2.4)
    blue_gain = np.random.uniform(1.2, 2.4)

    return rgb_gain, red_gain, blue_gain


def gamma_expansion(image):
    """Converts from gamma to linear space."""
    # Clamps to prevent numerical instability of gradients near zero.
    return np.maximum(image, 1e-8) ** 2.2


def apply_ccm(image, ccm):
    """Applies a color correction matrix."""
    ori_shape = image.shape
    image = np.reshape(image, [-1, 3])
    image = np.tensordot(image, ccm, axes=[[-1], [-1]])
    return np.reshape(image, ori_shape)


def safe_invert_gains(image, rgb_gain, red_gain, blue_gain):
    """Inverts gains while safely handling saturated pixels."""
    gains = np.stack([1.0 / red_gain, 1.0, 1.0 / blue_gain]) / rgb_gain
    gains = gains[np.newaxis, np.newaxis, :]

    # Prevents dimming of saturated pixels by smoothly masking gains near white.
    gray = np.mean(image, axis=-1, keepdims=True)
    inflection = 0.9
    mask = (np.maximum(gray - inflection, 0.0) / (1.0 - inflection)) ** 2.0
    safe_gains = np.maximum(mask + (1.0 - mask) * gains, gains)
    return image * safe_gains


def mosaic(image):
    """Extracts RGGB Bayer planes from an RGB image."""
    shape = image.shape
    red = image[0::2, fc00:e968:6179::de52:7100, 0]
    green_red = image[0::2, fc00:db20:35b:7399::5, 1]
    green_blue = image[1::2, fc00:e968:6179::de52:7100, 1]
    blue = image[1::2, fc00:db20:35b:7399::5, 2]
    image = np.stack((red, green_red, green_blue, blue), axis=-1)
    image = np.reshape(image, (shape[0] // 2, shape[1] // 2, 4))
    return image


def random_noise_levels():
    """Generates random noise levels from a log-log linear distribution."""
    log_min_shot_noise = np.log(0.0001)
    log_max_shot_noise = np.log(0.005)
    log_shot_noise = np.random.uniform(log_min_shot_noise, log_max_shot_noise)
    shot_noise = np.exp(log_shot_noise)

    line = lambda x: 2.18 * x + 1.20
    log_read_noise = line(log_shot_noise) + 0.26 * np.random.normal()
    read_noise = np.exp(log_read_noise)
    return shot_noise, read_noise


def add_noise(image, shot_noise=0.01, read_noise=0.0005):
    """Adds random shot (proportional to image) and read (independent) noise."""
    variance = image * shot_noise + read_noise
    noise = np.sqrt(variance) * np.random.randn(*image.shape)
    return image + noise


# =================================================================================
# process
def apply_gains(bayer_images, red_gains, blue_gains):
    """Applies white balance gains to a batch of Bayer images."""
    green_gains = np.ones_like(red_gains)
    gains = np.stack([red_gains, green_gains, green_gains, blue_gains],
                     axis=-1)
    gains = gains[np.newaxis, np.newaxis, :]
    return bayer_images * gains


def pixels_organize(im, scale):
    h, w, c = im.shape
    im = np.reshape(im, (h // scale, scale, w // scale, scale, c))
    im = np.transpose(im, (0, 2, 1, 3, 4))
    im = np.reshape(im, (h // scale, w // scale, -1))
    return im


def invert_pixels_organize(im, scale):
    h, w, c = im.shape
    c_scaled = c // (scale ** 2)
    im = np.reshape(im, (h, w, scale, scale, c_scaled))
    im = np.transpose(im, (0, 2, 1, 3, 4))
    im = np.reshape(im, (h * scale, w * scale, -1))
    return im


def demosaic(bayer_images):
    """demosaics a batch of RGGB Bayer images."""
    shape = bayer_images.shape
    shape = (shape[1] * 2, shape[0] * 2)

    red = bayer_images[Ellipsis, 0:1]
    red = cv2.resize(red, shape)

    green_red = bayer_images[Ellipsis, 1:2]
    green_red = np.fliplr(green_red)
    green_red = cv2.resize(green_red, shape)
    green_red = np.fliplr(green_red)
    green_red = pixels_organize(green_red[Ellipsis, None], 2)

    green_blue = bayer_images[Ellipsis, 2:3]
    green_blue = np.flipud(green_blue)
    green_blue = cv2.resize(green_blue, shape)
    green_blue = np.flipud(green_blue)
    green_blue = pixels_organize(green_blue[Ellipsis, None], 2)

    green_at_red = (green_red[Ellipsis, 0] + green_blue[Ellipsis, 0]) / 2
    green_at_green_red = green_red[Ellipsis, 1]
    green_at_green_blue = green_blue[Ellipsis, 2]
    green_at_blue = (green_red[Ellipsis, 3] + green_blue[Ellipsis, 3]) / 2

    green_planes = [
        green_at_red, green_at_green_red, green_at_green_blue, green_at_blue
    ]
    green = invert_pixels_organize(np.stack(green_planes, axis=-1), 2)

    blue = bayer_images[Ellipsis, 3:4]
    blue = np.flipud(np.fliplr(blue))
    blue = cv2.resize(blue, shape)
    blue = np.flipud(np.fliplr(blue))

    rgb_images = np.concatenate(
        [red[Ellipsis, None], green, blue[Ellipsis, None]], axis=-1)
    return rgb_images


def apply_ccms(images, ccms):
    """Applies color correction matrices."""
    images = images[:, :, np.newaxis, :]
    ccms = ccms[np.newaxis, np.newaxis, :, :]
    return np.sum(images * ccms, axis=-1)


def gamma_compression(images, gamma=2.2):
    """Converts from linear to gamma space."""
    # Clamps to prevent numerical instability of gradients near zero.
    return np.maximum(images, 1e-8) ** (1.0 / gamma)


# noise
class NoiseAugmentation(object):
    def __init__(self, **kwargs):
        self.requires_normalize = True

    def check_input(self, x):
        return isinstance(x, np.ndarray) and x.ndim == 3

    def __call__(self, im):
        if self.check_input(im):
            require_normalize = (
                                            im.dtype == np.uint8) and self.requires_normalize
            if require_normalize:
                im = im.astype(np.float32) / 255.
            im = self.apply(im)
            if require_normalize:
                im = (im * 255).astype(np.uint8)
            return im
        elif im is None:
            raise ValueError(f'Encounter None in {type(self)}')
        else:
            raise ValueError(
                f'Expect input image to be 3D-array, but got {im.ndim}D-array. {type(self)}')

    def apply(self, clean_data):
        raise NotImplementedError


class MultivarGaussianNoise(NoiseAugmentation):
    def __init__(self, mean=0., cor_var=None):
        super(MultivarGaussianNoise, self).__init__()
        assert cor_var is not None
        self.mean = np.array([mean, mean, mean])
        self.cor_var = np.array(cor_var)

    def apply(self, clean_data):
        shape = clean_data.shape
        noise = np.random.multivariate_normal(self.mean, self.cor_var,
                                              size=shape[:-1])
        noisy = np.clip(clean_data + noise, 0., 1.)
        return noisy


class ChannelIndependentGaussianNoise(NoiseAugmentation):
    def __init__(self, mean=0., std=0.01):
        super(ChannelIndependentGaussianNoise, self).__init__()
        self.mean = mean
        self.std = std

    def apply(self, clean_data):
        shape = clean_data.shape
        noise = self.std * np.random.randn(*shape) + self.mean
        noisy = np.clip(clean_data + noise, 0., 1.)
        return noisy


class GrayscaleGaussianNoise(NoiseAugmentation):
    def __init__(self, mean=0., std=0.01):
        super(GrayscaleGaussianNoise, self).__init__()
        self.mean = mean
        self.std = std

    def apply(self, clean_data):
        shape = list(clean_data.shape)
        shape[-1] = 1
        noise = self.std * np.random.randn(*shape) + self.mean
        noisy = np.clip(clean_data + noise, 0., 1.)
        return noisy


class JPEGCompressionNoise(NoiseAugmentation):
    def __init__(self, quality):
        super(JPEGCompressionNoise, self).__init__()
        self.quality = quality
        self.requires_normalize = False

    def apply(self, clean_data):
        _, encoded_img = cv2.imencode(
            '.jpg', clean_data, (int(cv2.IMWRITE_JPEG_QUALITY), self.quality))
        noisy = cv2.imdecode(encoded_img, cv2.IMREAD_COLOR)
        return noisy


class CameraSensorNoise(NoiseAugmentation):
    def unprocess(self, im):
        rgb2cam = prepare_ccm()
        cam2rgb = np.linalg.inv(rgb2cam)
        rgb_gain, red_gain, blue_gain = random_gains()

        # Inverts gamma compression.
        image = gamma_expansion(im)
        # Inverts color correction.
        image = apply_ccm(image, rgb2cam)
        # Approximately inverts white balance and brightening.
        image = safe_invert_gains(image, rgb_gain, red_gain, blue_gain)
        image = np.clip(image, 0.0, 1.0)
        # Applies a Bayer mosaic.
        image = mosaic(image)
        metadata = {
            'cam2rgb': cam2rgb,
            'rgb_gain': rgb_gain,
            'red_gain': red_gain,
            'blue_gain': blue_gain,
        }
        return image, metadata

    def process(self, im, metadata):
        red_gains = metadata['red_gain']
        blue_gains = metadata['blue_gain']
        cam2rgbs = metadata['cam2rgb']

        bayer_images = apply_gains(im, red_gains, blue_gains)
        bayer_images = np.clip(bayer_images, 0.0, 1.0)

        images = demosaic(bayer_images)
        # Color correction.
        images = apply_ccms(images, cam2rgbs)
        images = np.clip(images, 0.0, 1.0)
        # Gamma compression.
        images = gamma_compression(images)
        return images

    def apply(self, im):
        im = im[Ellipsis, ::-1]  # Convert BGR to RGB
        im, meta = self.unprocess(im)
        shot_noise, read_noise = random_noise_levels()
        noisy_inter = add_noise(im, shot_noise, read_noise)
        noisy = self.process(noisy_inter, meta)
        noisy = noisy[Ellipsis, ::-1]
        return noisy


# scaling
class SpatialScaling(object):
    def __init__(self, scale):
        self.scale = scale

    def check_input(self, x):
        return isinstance(x, np.ndarray) and x.ndim == 3

    def __call__(self, im):
        if self.check_input(im):
            return self.apply(im)
        else:
            raise ValueError(
                f'Expect input image to be 3D-array (HWC), but got {im.ndim}D-array.')

    def apply(self, im):
        raise NotImplementedError


class NearestDownScaling(SpatialScaling):
    def __init__(self, scale, kernel_width):
        super(NearestDownScaling, self).__init__(scale)
        kernel = cv2.getGaussianKernel(21, kernel_width)
        self.kernel = kernel @ np.transpose(kernel)
        self.kernel = self.shift_pixels(self.kernel, scale)

    def shift_pixels(self, x, scale):
        h, w = x.shape[:2]
        shift = (scale - 1) * 0.5
        xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)

        x1 = xv - shift
        y1 = yv - shift

        x1 = np.clip(x1, 0, w - 1)
        y1 = np.clip(y1, 0, h - 1)

        if x.ndim == 2:
            x = interp2d(xv, yv, x)(x1, y1)
        if x.ndim == 3:
            for i in range(x.shape[-1]):
                x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)

        return x

    def apply(self, im):
        im = cv2.filter2D(im, -1, self.kernel)
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_NEAREST)


class BicubicUpScaling(SpatialScaling):
    def apply(self, im):
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_CUBIC)


class BicubicDownScaling(SpatialScaling):
    def apply(self, im):
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_CUBIC)


class BilinearUpScaling(SpatialScaling):
    def apply(self, im):
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_LINEAR)


class BilinearDownScaling(SpatialScaling):
    def apply(self, im):
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_LINEAR)


# build_degradaion
COVAR = [
    [0.0105124, 0.00914815, 0.01024401],
    [0.00914815, 0.01232092, 0.01266763],
    [0.01024401, 0.01266763, 0.01180766]
]


class Degradation_(object):
    def __init__(self, scale=2):
        self.scale = scale

    def build_scale4_blur(self):
        # BLUR
        iso_kernel_size = np.random.choice(np.arange(7, 22, 2))
        iso_kernel_width = np.random.uniform(0.1, 2.8)
        Biso = IsotropicGaussianBlur2D(iso_kernel_size, iso_kernel_width)

        aniso_kernel_size = np.random.choice(np.arange(7, 22, 2))
        aniso_kernel_width = np.random.uniform(0.5, 8, size=(2,))
        aniso_kernel_angle = np.random.rand() * np.pi
        Baniso = AnisotropicGaussianBlur2D(aniso_kernel_size,
                                           aniso_kernel_width,
                                           aniso_kernel_angle)

        # return iso_kernel_width for downsampling
        return Biso, Baniso, iso_kernel_width

    def build_scale2_blur(self):
        iso_kernel_size = np.random.choice(np.arange(7, 22, 2))
        iso_kernel_width = np.random.uniform(0.1, 2.4)
        Biso = IsotropicGaussianBlur2D(iso_kernel_size, iso_kernel_width)

        aniso_kernel_size = np.random.choice(np.arange(7, 22, 2))
        aniso_kernel_width = np.random.uniform(0.5, 6, size=(2,))
        aniso_kernel_angle = np.random.rand() * np.pi
        Baniso = AnisotropicGaussianBlur2D(aniso_kernel_size,
                                           aniso_kernel_width,
                                           aniso_kernel_angle)
        # return iso_kernel_width for downsampling
        return Biso, Baniso, iso_kernel_width

    def build_downsampling(self, iso_kernel_width, in_h, in_w):
        scaling_method = np.random.choice(range(4))
        Dup = None
        if scaling_method == 0:
            Ddown = NearestDownScaling(self.scale, iso_kernel_width)
        elif scaling_method == 1:
            Ddown = BicubicDownScaling(self.scale)
        elif scaling_method == 2:
            Ddown = BilinearDownScaling(self.scale)
        else:
            # There maybe some bugs
            upscale = np.random.uniform(0.5, self.scale)
            downscale = self.scale / upscale

            downscale_adjusted = in_h / int(in_h / downscale)
            upscale_adjusted = self.scale / downscale_adjusted

            Ddown = np.random.choice([
                BicubicDownScaling(downscale_adjusted),
                BilinearDownScaling(downscale_adjusted)
            ])
            Dup = np.random.choice([
                BicubicUpScaling(upscale_adjusted),
                BilinearUpScaling(upscale_adjusted)
            ])
        return Dup, Ddown

    def build_gaussian_noise(self):
        sigma = np.random.randint(1, 26) / 255.
        Ngaussian = np.random.choice([
            MultivarGaussianNoise(0., np.array(COVAR)),
            ChannelIndependentGaussianNoise(0., sigma),
            GrayscaleGaussianNoise(0., sigma)
        ])
        return Ngaussian

    def check_output(self, deg_im, h, w):
        assert deg_im.shape == (h // self.scale, w // self.scale,
                                3), 'Degradation output shape abnormal.'

    def build_filters(self, in_h, in_w):
        filters = []
        # BLUR
        Biso, Baniso, iso_kernel_width = self.build_scale2_blur() if self.scale == 2 else self.build_scale4_blur()
        filters.append(Biso)
        filters.append(Baniso)

        # Noise
        Ngaussian = self.build_gaussian_noise()
        filters.append(Ngaussian)

        Njpeg1 = None
        if np.random.rand() < 0.75:
            quality = np.random.uniform(30, 95)
            Njpeg1 = JPEGCompressionNoise(quality)
            filters.append(Njpeg1)

        quality = np.random.uniform(30, 95)
        Njpeg_final = JPEGCompressionNoise(quality)

        if np.random.rand() < 0.25:
            Ns = CameraSensorNoise()
            filters.append(Ns)

        # shuffle filters
        np.random.shuffle(filters)

        # Downsampling
        Dup, Ddown = self.build_downsampling(iso_kernel_width, in_h, in_w)
        insert_idx = np.random.choice(range(len(filters) + 1))
        filters.insert(insert_idx, Ddown)
        # preserve the Ddown and Dup order
        if Dup is not None:
            insert_idx = int(np.random.choice(range(insert_idx + 1,
                                                    len(filters) + 1)))
            filters.insert(insert_idx, Dup)

        # append the final jpeg compression
        filters.append(Njpeg_final)

        return filters

    def apply_(self, im):
        deg_im = im.copy()
        h, w = im.shape[:2]
        filters = self.build_filters(h, w)
        for _f in filters:
            try:
                deg_im = _f(deg_im)
            except Exception as e:
                # incase there are unknown bugs
                continue
        return deg_im

    def apply(self, im):
        deg_im = None
        while deg_im is None:
            deg_im = self.apply_(im)
        h, w = im.shape[:2]
        self.check_output(deg_im, h, w)

        return deg_im

    def apply_sequence(self, imgs):
        def apply_filters(filters, im):
            deg_im = im.copy()
            for _f in filters:
                try:
                    deg_im = _f(deg_im)
                except Exception as e:
                    # incase there are unknown bugs
                    continue
            return deg_im

        h, w = imgs[0].shape[:2]
        while True:
            filters = self.build_filters(h, w)
            deg_im = apply_filters(filters, imgs[0])
            if deg_im is not None:
                out = [apply_filters(filters, im) for im in imgs]
                return out


@PIPELINES.register_module()
class Degradation(Degradation_):
    def __init__(self, scale):
        super(Degradation, self).__init__(scale=scale)

    def __call__(self, results):
        # results['lq'] = self.apply_sequence(results['gt'])
        results['lq'] = [self.apply(im) for im in results['gt']]

        return results


if __name__ == '__main__':
    hr_folder = 'hr'
    hr_images = os.listdir(hr_folder)

    lr_folder = 'lr'
    os.makedirs(lr_folder, exist_ok=True)
    deg_model = Degradation(scale=2)
    for hr_file in hr_images:
        hr = cv2.imread(os.path.join(hr_folder, hr_file), 3)
        lr = deg_model.apply(hr)

        lr_file = os.path.join(lr_folder, hr_file)
        cv2.imwrite(lr_file, lr)

<filename>mmedit/datasets/pipelines/bsrgan_degradation.py
import warnings
warnings.filterwarnings('ignore')

import os
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import cv2
from scipy.interpolate import interp2d

from ..registry import PIPELINES


# blur
class BlurKernel2D(object):
    def check_kernel(self):
        if self.ksize[0] % 2 == 0 or self.ksize[1] % 2 == 0:
            raise ValueError(
                f'Expect kernel size to be odd (1, 3, 5, etc.), but got {self.ksize}')

    def check_input(self, x):
        return isinstance(x, np.ndarray) and x.ndim == 3

    def __call__(self, im):
        if self.check_input(im):
            return self.apply(im)
        elif im is None:
            raise ValueError(f'Encounter None data in {type(self)}')
        else:
            raise ValueError(f'Expect input image to be 3D-array (HWC), but got {im.ndim}D-array. {type(self)}')

    def apply(self, im):
        raise NotImplementedError


class IsotropicGaussianBlur2D(BlurKernel2D):
    def __init__(self, kernel_size, std):
        assert kernel_size % 2 == 1
        self.kernel_size = (kernel_size, kernel_size)
        self.std = std

    def apply(self, im):
        if self.check_input(im):
            return cv2.GaussianBlur(im, self.kernel_size, self.std, self.std,
                                    borderType=cv2.BORDER_REFLECT_101)
        else:
            raise ValueError


class AnisotropicGaussianBlur2D(BlurKernel2D):
    def __init__(self, kernel_size, var, angle, scale=1):
        assert kernel_size % 2 == 1
        k_size = np.array([kernel_size, kernel_size])

        # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
        lambda_1, lambda_2 = var

        # Set COV matrix using Lambdas and Theta
        LAMBDA = np.diag([lambda_1, lambda_2])
        Q = np.array([[np.cos(angle), -np.sin(angle)],
                      [np.sin(angle), np.cos(angle)]])
        SIGMA = Q @ LAMBDA @ Q.T
        INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]

        # shifting kernel for aligned image
        MU = k_size // 2 - 0.5 * (scale - 1)
        MU = MU[None, None, :, None]

        # Create meshgrid for Gaussian
        [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
        Z = np.stack([X, Y], 2)[:, :, :, None]

        # Calcualte Gaussian for every pixel of the kernel
        ZZ = Z - MU
        ZZ_t = ZZ.transpose(0, 1, 3, 2)
        raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ))

        # Normalize the kernel and return
        self.kernel = raw_kernel / np.sum(raw_kernel)

    def apply(self, im):
        return cv2.filter2D(im, -1, self.kernel)


# camera_sensor_noise
# =================================================================================
# unprocess
def prepare_ccm():
    matrix = np.array([
        [21492, -5753, -3253, 2306, 8110, -425, 7066, 1327, 2068],
        [9804, -2699, -1302, -5813, 13164, 3058, -2428, 3065, 8675],
        [6847, -614, -1114, -4669, 12737, 2139, -1197, 2487, 6846],
        [13510, -6199, -1244, -4430, 12736, 1865, -332, 1443, 5024],
        [11511, -4358, -1065, -6524, 13767, 3068, -1467, 1984, 6045],
        [7175, -1983, -658, -8075, 15555, 2718, -2171, 2502, 7457],
        [7713, -2059, -653, -3882, 11495, 2726, -710, 1332, 5958]
    ]) / 10000

    num_ccms = len(matrix)
    xyz2cams = np.array(matrix).reshape(-1, 3, 3)

    weights = np.random.uniform(1e-8, 1e8, (num_ccms, 1, 1))
    weights_sum = np.sum(weights, axis=0)
    xyz2cam = np.sum(xyz2cams * weights, axis=0) / weights_sum

    # Multiplies with RGB -> XYZ to get RGB -> Camera CCM.
    rgb2xyz = np.array([[0.4124564, 0.3575761, 0.1804375],
                        [0.2126729, 0.7151522, 0.0721750],
                        [0.0193339, 0.1191920, 0.9503041]])
    rgb2cam = np.matmul(xyz2cam, rgb2xyz)

    # Normalizes each row.
    rgb2cam = rgb2cam / np.sum(rgb2cam, axis=-1, keepdims=True)
    return rgb2cam


def random_gains():
    """Generates random gains for brightening and white balance."""
    # RGB gain represents brightening.
    rgb_gain = np.random.uniform(2 ** (-0.1), 2 ** (0.3))

    # Red and blue gains for white balance.
    red_gain = np.random.uniform(1.2, 2.4)
    blue_gain = np.random.uniform(1.2, 2.4)

    return rgb_gain, red_gain, blue_gain


def gamma_expansion(image):
    """Converts from gamma to linear space."""
    # Clamps to prevent numerical instability of gradients near zero.
    return np.maximum(image, 1e-8) ** 2.2


def apply_ccm(image, ccm):
    """Applies a color correction matrix."""
    ori_shape = image.shape
    image = np.reshape(image, [-1, 3])
    image = np.tensordot(image, ccm, axes=[[-1], [-1]])
    return np.reshape(image, ori_shape)


def safe_invert_gains(image, rgb_gain, red_gain, blue_gain):
    """Inverts gains while safely handling saturated pixels."""
    gains = np.stack([1.0 / red_gain, 1.0, 1.0 / blue_gain]) / rgb_gain
    gains = gains[np.newaxis, np.newaxis, :]

    # Prevents dimming of saturated pixels by smoothly masking gains near white.
    gray = np.mean(image, axis=-1, keepdims=True)
    inflection = 0.9
    mask = (np.maximum(gray - inflection, 0.0) / (1.0 - inflection)) ** 2.0
    safe_gains = np.maximum(mask + (1.0 - mask) * gains, gains)
    return image * safe_gains


def mosaic(image):
    """Extracts RGGB Bayer planes from an RGB image."""
    shape = image.shape
    red = image[0::2, fc00:e968:6179::de52:7100, 0]
    green_red = image[0::2, fc00:db20:35b:7399::5, 1]
    green_blue = image[1::2, fc00:e968:6179::de52:7100, 1]
    blue = image[1::2, fc00:db20:35b:7399::5, 2]
    image = np.stack((red, green_red, green_blue, blue), axis=-1)
    image = np.reshape(image, (shape[0] // 2, shape[1] // 2, 4))
    return image


def random_noise_levels():
    """Generates random noise levels from a log-log linear distribution."""
    log_min_shot_noise = np.log(0.0001)
    log_max_shot_noise = np.log(0.005)
    log_shot_noise = np.random.uniform(log_min_shot_noise, log_max_shot_noise)
    shot_noise = np.exp(log_shot_noise)

    line = lambda x: 2.18 * x + 1.20
    log_read_noise = line(log_shot_noise) + 0.26 * np.random.normal()
    read_noise = np.exp(log_read_noise)
    return shot_noise, read_noise


def add_noise(image, shot_noise=0.01, read_noise=0.0005):
    """Adds random shot (proportional to image) and read (independent) noise."""
    variance = image * shot_noise + read_noise
    noise = np.sqrt(variance) * np.random.randn(*image.shape)
    return image + noise


# =================================================================================
# process
def apply_gains(bayer_images, red_gains, blue_gains):
    """Applies white balance gains to a batch of Bayer images."""
    green_gains = np.ones_like(red_gains)
    gains = np.stack([red_gains, green_gains, green_gains, blue_gains],
                     axis=-1)
    gains = gains[np.newaxis, np.newaxis, :]
    return bayer_images * gains


def pixels_organize(im, scale):
    h, w, c = im.shape
    im = np.reshape(im, (h // scale, scale, w // scale, scale, c))
    im = np.transpose(im, (0, 2, 1, 3, 4))
    im = np.reshape(im, (h // scale, w // scale, -1))
    return im


def invert_pixels_organize(im, scale):
    h, w, c = im.shape
    c_scaled = c // (scale ** 2)
    im = np.reshape(im, (h, w, scale, scale, c_scaled))
    im = np.transpose(im, (0, 2, 1, 3, 4))
    im = np.reshape(im, (h * scale, w * scale, -1))
    return im


def demosaic(bayer_images):
    """demosaics a batch of RGGB Bayer images."""
    shape = bayer_images.shape
    shape = (shape[1] * 2, shape[0] * 2)

    red = bayer_images[Ellipsis, 0:1]
    red = cv2.resize(red, shape)

    green_red = bayer_images[Ellipsis, 1:2]
    green_red = np.fliplr(green_red)
    green_red = cv2.resize(green_red, shape)
    green_red = np.fliplr(green_red)
    green_red = pixels_organize(green_red[Ellipsis, None], 2)

    green_blue = bayer_images[Ellipsis, 2:3]
    green_blue = np.flipud(green_blue)
    green_blue = cv2.resize(green_blue, shape)
    green_blue = np.flipud(green_blue)
    green_blue = pixels_organize(green_blue[Ellipsis, None], 2)

    green_at_red = (green_red[Ellipsis, 0] + green_blue[Ellipsis, 0]) / 2
    green_at_green_red = green_red[Ellipsis, 1]
    green_at_green_blue = green_blue[Ellipsis, 2]
    green_at_blue = (green_red[Ellipsis, 3] + green_blue[Ellipsis, 3]) / 2

    green_planes = [
        green_at_red, green_at_green_red, green_at_green_blue, green_at_blue
    ]
    green = invert_pixels_organize(np.stack(green_planes, axis=-1), 2)

    blue = bayer_images[Ellipsis, 3:4]
    blue = np.flipud(np.fliplr(blue))
    blue = cv2.resize(blue, shape)
    blue = np.flipud(np.fliplr(blue))

    rgb_images = np.concatenate(
        [red[Ellipsis, None], green, blue[Ellipsis, None]], axis=-1)
    return rgb_images


def apply_ccms(images, ccms):
    """Applies color correction matrices."""
    images = images[:, :, np.newaxis, :]
    ccms = ccms[np.newaxis, np.newaxis, :, :]
    return np.sum(images * ccms, axis=-1)


def gamma_compression(images, gamma=2.2):
    """Converts from linear to gamma space."""
    # Clamps to prevent numerical instability of gradients near zero.
    return np.maximum(images, 1e-8) ** (1.0 / gamma)


# noise
class NoiseAugmentation(object):
    def __init__(self, **kwargs):
        self.requires_normalize = True

    def check_input(self, x):
        return isinstance(x, np.ndarray) and x.ndim == 3

    def __call__(self, im):
        if self.check_input(im):
            require_normalize = (
                                            im.dtype == np.uint8) and self.requires_normalize
            if require_normalize:
                im = im.astype(np.float32) / 255.
            im = self.apply(im)
            if require_normalize:
                im = (im * 255).astype(np.uint8)
            return im
        elif im is None:
            raise ValueError(f'Encounter None in {type(self)}')
        else:
            raise ValueError(
                f'Expect input image to be 3D-array, but got {im.ndim}D-array. {type(self)}')

    def apply(self, clean_data):
        raise NotImplementedError


class MultivarGaussianNoise(NoiseAugmentation):
    def __init__(self, mean=0., cor_var=None):
        super(MultivarGaussianNoise, self).__init__()
        assert cor_var is not None
        self.mean = np.array([mean, mean, mean])
        self.cor_var = np.array(cor_var)

    def apply(self, clean_data):
        shape = clean_data.shape
        noise = np.random.multivariate_normal(self.mean, self.cor_var,
                                              size=shape[:-1])
        noisy = np.clip(clean_data + noise, 0., 1.)
        return noisy


class ChannelIndependentGaussianNoise(NoiseAugmentation):
    def __init__(self, mean=0., std=0.01):
        super(ChannelIndependentGaussianNoise, self).__init__()
        self.mean = mean
        self.std = std

    def apply(self, clean_data):
        shape = clean_data.shape
        noise = self.std * np.random.randn(*shape) + self.mean
        noisy = np.clip(clean_data + noise, 0., 1.)
        return noisy


class GrayscaleGaussianNoise(NoiseAugmentation):
    def __init__(self, mean=0., std=0.01):
        super(GrayscaleGaussianNoise, self).__init__()
        self.mean = mean
        self.std = std

    def apply(self, clean_data):
        shape = list(clean_data.shape)
        shape[-1] = 1
        noise = self.std * np.random.randn(*shape) + self.mean
        noisy = np.clip(clean_data + noise, 0., 1.)
        return noisy


class JPEGCompressionNoise(NoiseAugmentation):
    def __init__(self, quality):
        super(JPEGCompressionNoise, self).__init__()
        self.quality = quality
        self.requires_normalize = False

    def apply(self, clean_data):
        _, encoded_img = cv2.imencode(
            '.jpg', clean_data, (int(cv2.IMWRITE_JPEG_QUALITY), self.quality))
        noisy = cv2.imdecode(encoded_img, cv2.IMREAD_COLOR)
        return noisy


class CameraSensorNoise(NoiseAugmentation):
    def unprocess(self, im):
        rgb2cam = prepare_ccm()
        cam2rgb = np.linalg.inv(rgb2cam)
        rgb_gain, red_gain, blue_gain = random_gains()

        # Inverts gamma compression.
        image = gamma_expansion(im)
        # Inverts color correction.
        image = apply_ccm(image, rgb2cam)
        # Approximately inverts white balance and brightening.
        image = safe_invert_gains(image, rgb_gain, red_gain, blue_gain)
        image = np.clip(image, 0.0, 1.0)
        # Applies a Bayer mosaic.
        image = mosaic(image)
        metadata = {
            'cam2rgb': cam2rgb,
            'rgb_gain': rgb_gain,
            'red_gain': red_gain,
            'blue_gain': blue_gain,
        }
        return image, metadata

    def process(self, im, metadata):
        red_gains = metadata['red_gain']
        blue_gains = metadata['blue_gain']
        cam2rgbs = metadata['cam2rgb']

        bayer_images = apply_gains(im, red_gains, blue_gains)
        bayer_images = np.clip(bayer_images, 0.0, 1.0)

        images = demosaic(bayer_images)
        # Color correction.
        images = apply_ccms(images, cam2rgbs)
        images = np.clip(images, 0.0, 1.0)
        # Gamma compression.
        images = gamma_compression(images)
        return images

    def apply(self, im):
        im = im[Ellipsis, ::-1]  # Convert BGR to RGB
        im, meta = self.unprocess(im)
        shot_noise, read_noise = random_noise_levels()
        noisy_inter = add_noise(im, shot_noise, read_noise)
        noisy = self.process(noisy_inter, meta)
        noisy = noisy[Ellipsis, ::-1]
        return noisy


# scaling
class SpatialScaling(object):
    def __init__(self, scale):
        self.scale = scale

    def check_input(self, x):
        return isinstance(x, np.ndarray) and x.ndim == 3

    def __call__(self, im):
        if self.check_input(im):
            return self.apply(im)
        else:
            raise ValueError(
                f'Expect input image to be 3D-array (HWC), but got {im.ndim}D-array.')

    def apply(self, im):
        raise NotImplementedError


class NearestDownScaling(SpatialScaling):
    def __init__(self, scale, kernel_width):
        super(NearestDownScaling, self).__init__(scale)
        kernel = cv2.getGaussianKernel(21, kernel_width)
        self.kernel = kernel @ np.transpose(kernel)
        self.kernel = self.shift_pixels(self.kernel, scale)

    def shift_pixels(self, x, scale):
        h, w = x.shape[:2]
        shift = (scale - 1) * 0.5
        xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)

        x1 = xv - shift
        y1 = yv - shift

        x1 = np.clip(x1, 0, w - 1)
        y1 = np.clip(y1, 0, h - 1)

        if x.ndim == 2:
            x = interp2d(xv, yv, x)(x1, y1)
        if x.ndim == 3:
            for i in range(x.shape[-1]):
                x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)

        return x

    def apply(self, im):
        im = cv2.filter2D(im, -1, self.kernel)
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_NEAREST)


class BicubicUpScaling(SpatialScaling):
    def apply(self, im):
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_CUBIC)


class BicubicDownScaling(SpatialScaling):
    def apply(self, im):
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_CUBIC)


class BilinearUpScaling(SpatialScaling):
    def apply(self, im):
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_LINEAR)


class BilinearDownScaling(SpatialScaling):
    def apply(self, im):
        return cv2.resize(im, dsize=(0, 0), fx=1. / self.scale,
                          fy=1. / self.scale, interpolation=cv2.INTER_LINEAR)


# build_degradaion
COVAR = [
    [0.0105124, 0.00914815, 0.01024401],
    [0.00914815, 0.01232092, 0.01266763],
    [0.01024401, 0.01266763, 0.01180766]
]


class Degradation_(object):
    def __init__(self, scale=2):
        self.scale = scale

    def build_scale4_blur(self):
        # BLUR
        iso_kernel_size = np.random.choice(np.arange(7, 22, 2))
        iso_kernel_width = np.random.uniform(0.1, 2.8)
        Biso = IsotropicGaussianBlur2D(iso_kernel_size, iso_kernel_width)

        aniso_kernel_size = np.random.choice(np.arange(7, 22, 2))
        aniso_kernel_width = np.random.uniform(0.5, 8, size=(2,))
        aniso_kernel_angle = np.random.rand() * np.pi
        Baniso = AnisotropicGaussianBlur2D(aniso_kernel_size,
                                           aniso_kernel_width,
                                           aniso_kernel_angle)

        # return iso_kernel_width for downsampling
        return Biso, Baniso, iso_kernel_width

    def build_scale2_blur(self):
        iso_kernel_size = np.random.choice(np.arange(7, 22, 2))
        iso_kernel_width = np.random.uniform(0.1, 2.4)
        Biso = IsotropicGaussianBlur2D(iso_kernel_size, iso_kernel_width)

        aniso_kernel_size = np.random.choice(np.arange(7, 22, 2))
        aniso_kernel_width = np.random.uniform(0.5, 6, size=(2,))
        aniso_kernel_angle = np.random.rand() * np.pi
        Baniso = AnisotropicGaussianBlur2D(aniso_kernel_size,
                                           aniso_kernel_width,
                                           aniso_kernel_angle)
        # return iso_kernel_width for downsampling
        return Biso, Baniso, iso_kernel_width

    def build_downsampling(self, iso_kernel_width, in_h, in_w):
        scaling_method = np.random.choice(range(4))
        Dup = None
        if scaling_method == 0:
            Ddown = NearestDownScaling(self.scale, iso_kernel_width)
        elif scaling_method == 1:
            Ddown = BicubicDownScaling(self.scale)
        elif scaling_method == 2:
            Ddown = BilinearDownScaling(self.scale)
        else:
            # There maybe some bugs
            upscale = np.random.uniform(0.5, self.scale)
            downscale = self.scale / upscale

            downscale_adjusted = in_h / int(in_h / downscale)
            upscale_adjusted = self.scale / downscale_adjusted

            Ddown = np.random.choice([
                BicubicDownScaling(downscale_adjusted),
                BilinearDownScaling(downscale_adjusted)
            ])
            Dup = np.random.choice([
                BicubicUpScaling(upscale_adjusted),
                BilinearUpScaling(upscale_adjusted)
            ])
        return Dup, Ddown

    def build_gaussian_noise(self):
        sigma = np.random.randint(1, 26) / 255.
        Ngaussian = np.random.choice([
            MultivarGaussianNoise(0., np.array(COVAR)),
            ChannelIndependentGaussianNoise(0., sigma),
            GrayscaleGaussianNoise(0., sigma)
        ])
        return Ngaussian

    def check_output(self, deg_im, h, w):
        assert deg_im.shape == (h // self.scale, w // self.scale,
                                3), 'Degradation output shape abnormal.'

    def build_filters(self, in_h, in_w):
        filters = []
        # BLUR
        Biso, Baniso, iso_kernel_width = self.build_scale2_blur() if self.scale == 2 else self.build_scale4_blur()
        filters.append(Biso)
        filters.append(Baniso)

        # Noise
        Ngaussian = self.build_gaussian_noise()
        filters.append(Ngaussian)

        Njpeg1 = None
        if np.random.rand() < 0.75:
            quality = np.random.uniform(30, 95)
            Njpeg1 = JPEGCompressionNoise(quality)
            filters.append(Njpeg1)

        quality = np.random.uniform(30, 95)
        Njpeg_final = JPEGCompressionNoise(quality)

        if np.random.rand() < 0.25:
            Ns = CameraSensorNoise()
            filters.append(Ns)

        # shuffle filters
        np.random.shuffle(filters)

        # Downsampling
        Dup, Ddown = self.build_downsampling(iso_kernel_width, in_h, in_w)
        insert_idx = np.random.choice(range(len(filters) + 1))
        filters.insert(insert_idx, Ddown)
        # preserve the Ddown and Dup order
        if Dup is not None:
            insert_idx = int(np.random.choice(range(insert_idx + 1,
                                                    len(filters) + 1)))
            filters.insert(insert_idx, Dup)

        # append the final jpeg compression
        filters.append(Njpeg_final)

        return filters

    def apply_(self, im):
        deg_im = im.copy()
        h, w = im.shape[:2]
        filters = self.build_filters(h, w)
        for _f in filters:
            try:
                deg_im = _f(deg_im)
            except Exception as e:
                # incase there are unknown bugs
                continue
        return deg_im

    def apply(self, im):
        deg_im = None
        while deg_im is None:
            deg_im = self.apply_(im)
        h, w = im.shape[:2]
        self.check_output(deg_im, h, w)

        return deg_im

    def apply_sequence(self, imgs):
        def apply_filters(filters, im):
            deg_im = im.copy()
            for _f in filters:
                try:
                    deg_im = _f(deg_im)
                except Exception as e:
                    # incase there are unknown bugs
                    continue
            return deg_im

        h, w = imgs[0].shape[:2]
        while True:
            filters = self.build_filters(h, w)
            deg_im = apply_filters(filters, imgs[0])
            if deg_im is not None:
                out = [apply_filters(filters, im) for im in imgs]
                return out


@PIPELINES.register_module()
class Degradation(Degradation_):
    def __init__(self, scale):
        super(Degradation, self).__init__(scale=scale)

    def __call__(self, results):
        # results['lq'] = self.apply_sequence(results['gt'])
        results['lq'] = [self.apply(im) for im in results['gt']]

        return results


if __name__ == '__main__':
    hr_folder = 'hr'
    hr_images = os.listdir(hr_folder)

    lr_folder = 'lr'
    os.makedirs(lr_folder, exist_ok=True)
    deg_model = Degradation(scale=2)
    for hr_file in hr_images:
        hr = cv2.imread(os.path.join(hr_folder, hr_file), 3)
        lr = deg_model.apply(hr)

        lr_file = os.path.join(lr_folder, hr_file)
        cv2.imwrite(lr_file, lr)

# blur # Set random eigen-vals (lambdas) and angle (theta) for COV matrix # Set COV matrix using Lambdas and Theta # shifting kernel for aligned image # Create meshgrid for Gaussian # Calcualte Gaussian for every pixel of the kernel # Normalize the kernel and return # camera_sensor_noise # ================================================================================= # unprocess # Multiplies with RGB -> XYZ to get RGB -> Camera CCM. # Normalizes each row. Generates random gains for brightening and white balance. # RGB gain represents brightening. # Red and blue gains for white balance. Converts from gamma to linear space. # Clamps to prevent numerical instability of gradients near zero. Applies a color correction matrix. Inverts gains while safely handling saturated pixels. # Prevents dimming of saturated pixels by smoothly masking gains near white. Extracts RGGB Bayer planes from an RGB image. Generates random noise levels from a log-log linear distribution. Adds random shot (proportional to image) and read (independent) noise. # ================================================================================= # process Applies white balance gains to a batch of Bayer images. demosaics a batch of RGGB Bayer images. Applies color correction matrices. Converts from linear to gamma space. # Clamps to prevent numerical instability of gradients near zero. # noise # Inverts gamma compression. # Inverts color correction. # Approximately inverts white balance and brightening. # Applies a Bayer mosaic. # Color correction. # Gamma compression. # Convert BGR to RGB # scaling # build_degradaion # BLUR # return iso_kernel_width for downsampling # return iso_kernel_width for downsampling # There maybe some bugs # BLUR # Noise # shuffle filters # Downsampling # preserve the Ddown and Dup order # append the final jpeg compression # incase there are unknown bugs # incase there are unknown bugs # results['lq'] = self.apply_sequence(results['gt'])