Spaces:
Running
on
A10G
Running
on
A10G
| # -*- coding: utf-8 -*- | |
| """ | |
| # -------------------------------------------- | |
| # Super-Resolution | |
| # -------------------------------------------- | |
| # | |
| # Kai Zhang (cskaizhang@gmail.com) | |
| # https://github.com/cszn | |
| # From 2019/03--2021/08 | |
| # -------------------------------------------- | |
| """ | |
| import numpy as np | |
| import cv2 | |
| import torch | |
| from functools import partial | |
| import random | |
| from scipy import ndimage | |
| import scipy | |
| import scipy.stats as ss | |
| from scipy.interpolate import interp2d | |
| from scipy.linalg import orth | |
| import albumentations | |
| import ldm.modules.image_degradation.utils_image as util | |
| def modcrop_np(img, sf): | |
| ''' | |
| Args: | |
| img: numpy image, WxH or WxHxC | |
| sf: scale factor | |
| Return: | |
| cropped image | |
| ''' | |
| w, h = img.shape[:2] | |
| im = np.copy(img) | |
| return im[:w - w % sf, :h - h % sf, ...] | |
| """ | |
| # -------------------------------------------- | |
| # anisotropic Gaussian kernels | |
| # -------------------------------------------- | |
| """ | |
| def analytic_kernel(k): | |
| """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)""" | |
| k_size = k.shape[0] | |
| # Calculate the big kernels size | |
| big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2)) | |
| # Loop over the small kernel to fill the big one | |
| for r in range(k_size): | |
| for c in range(k_size): | |
| big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k | |
| # Crop the edges of the big kernel to ignore very small values and increase run time of SR | |
| crop = k_size // 2 | |
| cropped_big_k = big_k[crop:-crop, crop:-crop] | |
| # Normalize to 1 | |
| return cropped_big_k / cropped_big_k.sum() | |
| def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6): | |
| """ generate an anisotropic Gaussian kernel | |
| Args: | |
| ksize : e.g., 15, kernel size | |
| theta : [0, pi], rotation angle range | |
| l1 : [0.1,50], scaling of eigenvalues | |
| l2 : [0.1,l1], scaling of eigenvalues | |
| If l1 = l2, will get an isotropic Gaussian kernel. | |
| Returns: | |
| k : kernel | |
| """ | |
| v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.])) | |
| V = np.array([[v[0], v[1]], [v[1], -v[0]]]) | |
| D = np.array([[l1, 0], [0, l2]]) | |
| Sigma = np.dot(np.dot(V, D), np.linalg.inv(V)) | |
| k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize) | |
| return k | |
| def gm_blur_kernel(mean, cov, size=15): | |
| center = size / 2.0 + 0.5 | |
| k = np.zeros([size, size]) | |
| for y in range(size): | |
| for x in range(size): | |
| cy = y - center + 1 | |
| cx = x - center + 1 | |
| k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov) | |
| k = k / np.sum(k) | |
| return k | |
| def shift_pixel(x, sf, upper_left=True): | |
| """shift pixel for super-resolution with different scale factors | |
| Args: | |
| x: WxHxC or WxH | |
| sf: scale factor | |
| upper_left: shift direction | |
| """ | |
| h, w = x.shape[:2] | |
| shift = (sf - 1) * 0.5 | |
| xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0) | |
| if upper_left: | |
| x1 = xv + shift | |
| y1 = yv + shift | |
| else: | |
| 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 blur(x, k): | |
| ''' | |
| x: image, NxcxHxW | |
| k: kernel, Nx1xhxw | |
| ''' | |
| n, c = x.shape[:2] | |
| p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2 | |
| x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate') | |
| k = k.repeat(1, c, 1, 1) | |
| k = k.view(-1, 1, k.shape[2], k.shape[3]) | |
| x = x.view(1, -1, x.shape[2], x.shape[3]) | |
| x = torch.nn.functional.conv2d(x, k, bias=None, stride=1, padding=0, groups=n * c) | |
| x = x.view(n, c, x.shape[2], x.shape[3]) | |
| return x | |
| def gen_kernel(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0): | |
| """" | |
| # modified version of https://github.com/assafshocher/BlindSR_dataset_generator | |
| # Kai Zhang | |
| # min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var | |
| # max_var = 2.5 * sf | |
| """ | |
| # Set random eigen-vals (lambdas) and angle (theta) for COV matrix | |
| lambda_1 = min_var + np.random.rand() * (max_var - min_var) | |
| lambda_2 = min_var + np.random.rand() * (max_var - min_var) | |
| theta = np.random.rand() * np.pi # random theta | |
| noise = -noise_level + np.random.rand(*k_size) * noise_level * 2 | |
| # Set COV matrix using Lambdas and Theta | |
| LAMBDA = np.diag([lambda_1, lambda_2]) | |
| Q = np.array([[np.cos(theta), -np.sin(theta)], | |
| [np.sin(theta), np.cos(theta)]]) | |
| SIGMA = Q @ LAMBDA @ Q.T | |
| INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :] | |
| # Set expectation position (shifting kernel for aligned image) | |
| MU = k_size // 2 - 0.5 * (scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2) | |
| 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)) * (1 + noise) | |
| # shift the kernel so it will be centered | |
| # raw_kernel_centered = kernel_shift(raw_kernel, scale_factor) | |
| # Normalize the kernel and return | |
| # kernel = raw_kernel_centered / np.sum(raw_kernel_centered) | |
| kernel = raw_kernel / np.sum(raw_kernel) | |
| return kernel | |
| def fspecial_gaussian(hsize, sigma): | |
| hsize = [hsize, hsize] | |
| siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0] | |
| std = sigma | |
| [x, y] = np.meshgrid(np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1)) | |
| arg = -(x * x + y * y) / (2 * std * std) | |
| h = np.exp(arg) | |
| h[h < scipy.finfo(float).eps * h.max()] = 0 | |
| sumh = h.sum() | |
| if sumh != 0: | |
| h = h / sumh | |
| return h | |
| def fspecial_laplacian(alpha): | |
| alpha = max([0, min([alpha, 1])]) | |
| h1 = alpha / (alpha + 1) | |
| h2 = (1 - alpha) / (alpha + 1) | |
| h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]] | |
| h = np.array(h) | |
| return h | |
| def fspecial(filter_type, *args, **kwargs): | |
| ''' | |
| python code from: | |
| https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py | |
| ''' | |
| if filter_type == 'gaussian': | |
| return fspecial_gaussian(*args, **kwargs) | |
| if filter_type == 'laplacian': | |
| return fspecial_laplacian(*args, **kwargs) | |
| """ | |
| # -------------------------------------------- | |
| # degradation models | |
| # -------------------------------------------- | |
| """ | |
| def bicubic_degradation(x, sf=3): | |
| ''' | |
| Args: | |
| x: HxWxC image, [0, 1] | |
| sf: down-scale factor | |
| Return: | |
| bicubicly downsampled LR image | |
| ''' | |
| x = util.imresize_np(x, scale=1 / sf) | |
| return x | |
| def srmd_degradation(x, k, sf=3): | |
| ''' blur + bicubic downsampling | |
| Args: | |
| x: HxWxC image, [0, 1] | |
| k: hxw, double | |
| sf: down-scale factor | |
| Return: | |
| downsampled LR image | |
| Reference: | |
| @inproceedings{zhang2018learning, | |
| title={Learning a single convolutional super-resolution network for multiple degradations}, | |
| author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei}, | |
| booktitle={IEEE Conference on Computer Vision and Pattern Recognition}, | |
| pages={3262--3271}, | |
| year={2018} | |
| } | |
| ''' | |
| x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') # 'nearest' | 'mirror' | |
| x = bicubic_degradation(x, sf=sf) | |
| return x | |
| def dpsr_degradation(x, k, sf=3): | |
| ''' bicubic downsampling + blur | |
| Args: | |
| x: HxWxC image, [0, 1] | |
| k: hxw, double | |
| sf: down-scale factor | |
| Return: | |
| downsampled LR image | |
| Reference: | |
| @inproceedings{zhang2019deep, | |
| title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels}, | |
| author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei}, | |
| booktitle={IEEE Conference on Computer Vision and Pattern Recognition}, | |
| pages={1671--1681}, | |
| year={2019} | |
| } | |
| ''' | |
| x = bicubic_degradation(x, sf=sf) | |
| x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') | |
| return x | |
| def classical_degradation(x, k, sf=3): | |
| ''' blur + downsampling | |
| Args: | |
| x: HxWxC image, [0, 1]/[0, 255] | |
| k: hxw, double | |
| sf: down-scale factor | |
| Return: | |
| downsampled LR image | |
| ''' | |
| x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') | |
| # x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2)) | |
| st = 0 | |
| return x[st::sf, st::sf, ...] | |
| def add_sharpening(img, weight=0.5, radius=50, threshold=10): | |
| """USM sharpening. borrowed from real-ESRGAN | |
| Input image: I; Blurry image: B. | |
| 1. K = I + weight * (I - B) | |
| 2. Mask = 1 if abs(I - B) > threshold, else: 0 | |
| 3. Blur mask: | |
| 4. Out = Mask * K + (1 - Mask) * I | |
| Args: | |
| img (Numpy array): Input image, HWC, BGR; float32, [0, 1]. | |
| weight (float): Sharp weight. Default: 1. | |
| radius (float): Kernel size of Gaussian blur. Default: 50. | |
| threshold (int): | |
| """ | |
| if radius % 2 == 0: | |
| radius += 1 | |
| blur = cv2.GaussianBlur(img, (radius, radius), 0) | |
| residual = img - blur | |
| mask = np.abs(residual) * 255 > threshold | |
| mask = mask.astype('float32') | |
| soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0) | |
| K = img + weight * residual | |
| K = np.clip(K, 0, 1) | |
| return soft_mask * K + (1 - soft_mask) * img | |
| def add_blur(img, sf=4): | |
| wd2 = 4.0 + sf | |
| wd = 2.0 + 0.2 * sf | |
| if random.random() < 0.5: | |
| l1 = wd2 * random.random() | |
| l2 = wd2 * random.random() | |
| k = anisotropic_Gaussian(ksize=2 * random.randint(2, 11) + 3, theta=random.random() * np.pi, l1=l1, l2=l2) | |
| else: | |
| k = fspecial('gaussian', 2 * random.randint(2, 11) + 3, wd * random.random()) | |
| img = ndimage.filters.convolve(img, np.expand_dims(k, axis=2), mode='mirror') | |
| return img | |
| def add_resize(img, sf=4): | |
| rnum = np.random.rand() | |
| if rnum > 0.8: # up | |
| sf1 = random.uniform(1, 2) | |
| elif rnum < 0.7: # down | |
| sf1 = random.uniform(0.5 / sf, 1) | |
| else: | |
| sf1 = 1.0 | |
| img = cv2.resize(img, (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])), interpolation=random.choice([1, 2, 3])) | |
| img = np.clip(img, 0.0, 1.0) | |
| return img | |
| # def add_Gaussian_noise(img, noise_level1=2, noise_level2=25): | |
| # noise_level = random.randint(noise_level1, noise_level2) | |
| # rnum = np.random.rand() | |
| # if rnum > 0.6: # add color Gaussian noise | |
| # img += np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32) | |
| # elif rnum < 0.4: # add grayscale Gaussian noise | |
| # img += np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32) | |
| # else: # add noise | |
| # L = noise_level2 / 255. | |
| # D = np.diag(np.random.rand(3)) | |
| # U = orth(np.random.rand(3, 3)) | |
| # conv = np.dot(np.dot(np.transpose(U), D), U) | |
| # img += np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32) | |
| # img = np.clip(img, 0.0, 1.0) | |
| # return img | |
| def add_Gaussian_noise(img, noise_level1=2, noise_level2=25): | |
| noise_level = random.randint(noise_level1, noise_level2) | |
| rnum = np.random.rand() | |
| if rnum > 0.6: # add color Gaussian noise | |
| img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32) | |
| elif rnum < 0.4: # add grayscale Gaussian noise | |
| img = img + np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32) | |
| else: # add noise | |
| L = noise_level2 / 255. | |
| D = np.diag(np.random.rand(3)) | |
| U = orth(np.random.rand(3, 3)) | |
| conv = np.dot(np.dot(np.transpose(U), D), U) | |
| img = img + np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32) | |
| img = np.clip(img, 0.0, 1.0) | |
| return img | |
| def add_speckle_noise(img, noise_level1=2, noise_level2=25): | |
| noise_level = random.randint(noise_level1, noise_level2) | |
| img = np.clip(img, 0.0, 1.0) | |
| rnum = random.random() | |
| if rnum > 0.6: | |
| img += img * np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32) | |
| elif rnum < 0.4: | |
| img += img * np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32) | |
| else: | |
| L = noise_level2 / 255. | |
| D = np.diag(np.random.rand(3)) | |
| U = orth(np.random.rand(3, 3)) | |
| conv = np.dot(np.dot(np.transpose(U), D), U) | |
| img += img * np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32) | |
| img = np.clip(img, 0.0, 1.0) | |
| return img | |
| def add_Poisson_noise(img): | |
| img = np.clip((img * 255.0).round(), 0, 255) / 255. | |
| vals = 10 ** (2 * random.random() + 2.0) # [2, 4] | |
| if random.random() < 0.5: | |
| img = np.random.poisson(img * vals).astype(np.float32) / vals | |
| else: | |
| img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114]) | |
| img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255. | |
| noise_gray = np.random.poisson(img_gray * vals).astype(np.float32) / vals - img_gray | |
| img += noise_gray[:, :, np.newaxis] | |
| img = np.clip(img, 0.0, 1.0) | |
| return img | |
| def add_JPEG_noise(img): | |
| quality_factor = random.randint(30, 95) | |
| img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR) | |
| result, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor]) | |
| img = cv2.imdecode(encimg, 1) | |
| img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB) | |
| return img | |
| def random_crop(lq, hq, sf=4, lq_patchsize=64): | |
| h, w = lq.shape[:2] | |
| rnd_h = random.randint(0, h - lq_patchsize) | |
| rnd_w = random.randint(0, w - lq_patchsize) | |
| lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :] | |
| rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf) | |
| hq = hq[rnd_h_H:rnd_h_H + lq_patchsize * sf, rnd_w_H:rnd_w_H + lq_patchsize * sf, :] | |
| return lq, hq | |
| def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None): | |
| """ | |
| This is the degradation model of BSRGAN from the paper | |
| "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution" | |
| ---------- | |
| img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf) | |
| sf: scale factor | |
| isp_model: camera ISP model | |
| Returns | |
| ------- | |
| img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1] | |
| hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1] | |
| """ | |
| isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25 | |
| sf_ori = sf | |
| h1, w1 = img.shape[:2] | |
| img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop | |
| h, w = img.shape[:2] | |
| if h < lq_patchsize * sf or w < lq_patchsize * sf: | |
| raise ValueError(f'img size ({h1}X{w1}) is too small!') | |
| hq = img.copy() | |
| if sf == 4 and random.random() < scale2_prob: # downsample1 | |
| if np.random.rand() < 0.5: | |
| img = cv2.resize(img, (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])), | |
| interpolation=random.choice([1, 2, 3])) | |
| else: | |
| img = util.imresize_np(img, 1 / 2, True) | |
| img = np.clip(img, 0.0, 1.0) | |
| sf = 2 | |
| shuffle_order = random.sample(range(7), 7) | |
| idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3) | |
| if idx1 > idx2: # keep downsample3 last | |
| shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1] | |
| for i in shuffle_order: | |
| if i == 0: | |
| img = add_blur(img, sf=sf) | |
| elif i == 1: | |
| img = add_blur(img, sf=sf) | |
| elif i == 2: | |
| a, b = img.shape[1], img.shape[0] | |
| # downsample2 | |
| if random.random() < 0.75: | |
| sf1 = random.uniform(1, 2 * sf) | |
| img = cv2.resize(img, (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])), | |
| interpolation=random.choice([1, 2, 3])) | |
| else: | |
| k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf)) | |
| k_shifted = shift_pixel(k, sf) | |
| k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel | |
| img = ndimage.filters.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror') | |
| img = img[0::sf, 0::sf, ...] # nearest downsampling | |
| img = np.clip(img, 0.0, 1.0) | |
| elif i == 3: | |
| # downsample3 | |
| img = cv2.resize(img, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3])) | |
| img = np.clip(img, 0.0, 1.0) | |
| elif i == 4: | |
| # add Gaussian noise | |
| img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25) | |
| elif i == 5: | |
| # add JPEG noise | |
| if random.random() < jpeg_prob: | |
| img = add_JPEG_noise(img) | |
| elif i == 6: | |
| # add processed camera sensor noise | |
| if random.random() < isp_prob and isp_model is not None: | |
| with torch.no_grad(): | |
| img, hq = isp_model.forward(img.copy(), hq) | |
| # add final JPEG compression noise | |
| img = add_JPEG_noise(img) | |
| # random crop | |
| img, hq = random_crop(img, hq, sf_ori, lq_patchsize) | |
| return img, hq | |
| # todo no isp_model? | |
| def degradation_bsrgan_variant(image, sf=4, isp_model=None): | |
| """ | |
| This is the degradation model of BSRGAN from the paper | |
| "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution" | |
| ---------- | |
| sf: scale factor | |
| isp_model: camera ISP model | |
| Returns | |
| ------- | |
| img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1] | |
| hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1] | |
| """ | |
| image = util.uint2single(image) | |
| isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25 | |
| sf_ori = sf | |
| h1, w1 = image.shape[:2] | |
| image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop | |
| h, w = image.shape[:2] | |
| hq = image.copy() | |
| if sf == 4 and random.random() < scale2_prob: # downsample1 | |
| if np.random.rand() < 0.5: | |
| image = cv2.resize(image, (int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])), | |
| interpolation=random.choice([1, 2, 3])) | |
| else: | |
| image = util.imresize_np(image, 1 / 2, True) | |
| image = np.clip(image, 0.0, 1.0) | |
| sf = 2 | |
| shuffle_order = random.sample(range(7), 7) | |
| idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3) | |
| if idx1 > idx2: # keep downsample3 last | |
| shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1] | |
| for i in shuffle_order: | |
| if i == 0: | |
| image = add_blur(image, sf=sf) | |
| elif i == 1: | |
| image = add_blur(image, sf=sf) | |
| elif i == 2: | |
| a, b = image.shape[1], image.shape[0] | |
| # downsample2 | |
| if random.random() < 0.75: | |
| sf1 = random.uniform(1, 2 * sf) | |
| image = cv2.resize(image, (int(1 / sf1 * image.shape[1]), int(1 / sf1 * image.shape[0])), | |
| interpolation=random.choice([1, 2, 3])) | |
| else: | |
| k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf)) | |
| k_shifted = shift_pixel(k, sf) | |
| k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel | |
| image = ndimage.filters.convolve(image, np.expand_dims(k_shifted, axis=2), mode='mirror') | |
| image = image[0::sf, 0::sf, ...] # nearest downsampling | |
| image = np.clip(image, 0.0, 1.0) | |
| elif i == 3: | |
| # downsample3 | |
| image = cv2.resize(image, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3])) | |
| image = np.clip(image, 0.0, 1.0) | |
| elif i == 4: | |
| # add Gaussian noise | |
| image = add_Gaussian_noise(image, noise_level1=2, noise_level2=25) | |
| elif i == 5: | |
| # add JPEG noise | |
| if random.random() < jpeg_prob: | |
| image = add_JPEG_noise(image) | |
| # elif i == 6: | |
| # # add processed camera sensor noise | |
| # if random.random() < isp_prob and isp_model is not None: | |
| # with torch.no_grad(): | |
| # img, hq = isp_model.forward(img.copy(), hq) | |
| # add final JPEG compression noise | |
| image = add_JPEG_noise(image) | |
| image = util.single2uint(image) | |
| example = {"image":image} | |
| return example | |
| # TODO incase there is a pickle error one needs to replace a += x with a = a + x in add_speckle_noise etc... | |
| def degradation_bsrgan_plus(img, sf=4, shuffle_prob=0.5, use_sharp=True, lq_patchsize=64, isp_model=None): | |
| """ | |
| This is an extended degradation model by combining | |
| the degradation models of BSRGAN and Real-ESRGAN | |
| ---------- | |
| img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf) | |
| sf: scale factor | |
| use_shuffle: the degradation shuffle | |
| use_sharp: sharpening the img | |
| Returns | |
| ------- | |
| img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1] | |
| hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1] | |
| """ | |
| h1, w1 = img.shape[:2] | |
| img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop | |
| h, w = img.shape[:2] | |
| if h < lq_patchsize * sf or w < lq_patchsize * sf: | |
| raise ValueError(f'img size ({h1}X{w1}) is too small!') | |
| if use_sharp: | |
| img = add_sharpening(img) | |
| hq = img.copy() | |
| if random.random() < shuffle_prob: | |
| shuffle_order = random.sample(range(13), 13) | |
| else: | |
| shuffle_order = list(range(13)) | |
| # local shuffle for noise, JPEG is always the last one | |
| shuffle_order[2:6] = random.sample(shuffle_order[2:6], len(range(2, 6))) | |
| shuffle_order[9:13] = random.sample(shuffle_order[9:13], len(range(9, 13))) | |
| poisson_prob, speckle_prob, isp_prob = 0.1, 0.1, 0.1 | |
| for i in shuffle_order: | |
| if i == 0: | |
| img = add_blur(img, sf=sf) | |
| elif i == 1: | |
| img = add_resize(img, sf=sf) | |
| elif i == 2: | |
| img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25) | |
| elif i == 3: | |
| if random.random() < poisson_prob: | |
| img = add_Poisson_noise(img) | |
| elif i == 4: | |
| if random.random() < speckle_prob: | |
| img = add_speckle_noise(img) | |
| elif i == 5: | |
| if random.random() < isp_prob and isp_model is not None: | |
| with torch.no_grad(): | |
| img, hq = isp_model.forward(img.copy(), hq) | |
| elif i == 6: | |
| img = add_JPEG_noise(img) | |
| elif i == 7: | |
| img = add_blur(img, sf=sf) | |
| elif i == 8: | |
| img = add_resize(img, sf=sf) | |
| elif i == 9: | |
| img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25) | |
| elif i == 10: | |
| if random.random() < poisson_prob: | |
| img = add_Poisson_noise(img) | |
| elif i == 11: | |
| if random.random() < speckle_prob: | |
| img = add_speckle_noise(img) | |
| elif i == 12: | |
| if random.random() < isp_prob and isp_model is not None: | |
| with torch.no_grad(): | |
| img, hq = isp_model.forward(img.copy(), hq) | |
| else: | |
| print('check the shuffle!') | |
| # resize to desired size | |
| img = cv2.resize(img, (int(1 / sf * hq.shape[1]), int(1 / sf * hq.shape[0])), | |
| interpolation=random.choice([1, 2, 3])) | |
| # add final JPEG compression noise | |
| img = add_JPEG_noise(img) | |
| # random crop | |
| img, hq = random_crop(img, hq, sf, lq_patchsize) | |
| return img, hq | |
| if __name__ == '__main__': | |
| print("hey") | |
| img = util.imread_uint('utils/test.png', 3) | |
| print(img) | |
| img = util.uint2single(img) | |
| print(img) | |
| img = img[:448, :448] | |
| h = img.shape[0] // 4 | |
| print("resizing to", h) | |
| sf = 4 | |
| deg_fn = partial(degradation_bsrgan_variant, sf=sf) | |
| for i in range(20): | |
| print(i) | |
| img_lq = deg_fn(img) | |
| print(img_lq) | |
| img_lq_bicubic = albumentations.SmallestMaxSize(max_size=h, interpolation=cv2.INTER_CUBIC)(image=img)["image"] | |
| print(img_lq.shape) | |
| print("bicubic", img_lq_bicubic.shape) | |
| print(img_hq.shape) | |
| lq_nearest = cv2.resize(util.single2uint(img_lq), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])), | |
| interpolation=0) | |
| lq_bicubic_nearest = cv2.resize(util.single2uint(img_lq_bicubic), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])), | |
| interpolation=0) | |
| img_concat = np.concatenate([lq_bicubic_nearest, lq_nearest, util.single2uint(img_hq)], axis=1) | |
| util.imsave(img_concat, str(i) + '.png') | |