import cv2 import math import numpy as np import random import torch from torch.utils import data as data from basicsr.data.degradations import circular_lowpass_kernel, random_mixed_kernels from basicsr.data.transforms import augment from basicsr.utils import img2tensor, DiffJPEG, USMSharp from basicsr.utils.img_process_util import filter2D from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt from basicsr.data.transforms import paired_random_crop AUGMENT_OPT = { 'use_hflip': False, 'use_rot': False } KERNEL_OPT = { 'blur_kernel_size': 21, 'kernel_list': ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso'], 'kernel_prob': [0.45, 0.25, 0.12, 0.03, 0.12, 0.03], 'sinc_prob': 0.1, 'blur_sigma': [0.2, 3], 'betag_range': [0.5, 4], 'betap_range': [1, 2], 'blur_kernel_size2': 21, 'kernel_list2': ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso'], 'kernel_prob2': [0.45, 0.25, 0.12, 0.03, 0.12, 0.03], 'sinc_prob2': 0.1, 'blur_sigma2': [0.2, 1.5], 'betag_range2': [0.5, 4], 'betap_range2': [1, 2], 'final_sinc_prob': 0.8, } DEGRADE_OPT = { 'resize_prob': [0.2, 0.7, 0.1], # up, down, keep 'resize_range': [0.15, 1.5], 'gaussian_noise_prob': 0.5, 'noise_range': [1, 30], 'poisson_scale_range': [0.05, 3], 'gray_noise_prob': 0.4, 'jpeg_range': [30, 95], # the second degradation process 'second_blur_prob': 0.8, 'resize_prob2': [0.3, 0.4, 0.3], # up, down, keep 'resize_range2': [0.3, 1.2], 'gaussian_noise_prob2': 0.5, 'noise_range2': [1, 25], 'poisson_scale_range2': [0.05, 2.5], 'gray_noise_prob2': 0.4, 'jpeg_range2': [30, 95], 'gt_size': 512, 'no_degradation_prob': 0.01, 'use_usm': True, 'sf': 4, 'random_size': False, 'resize_lq': True } class RealESRGANDegradation: def __init__(self, augment_opt=None, kernel_opt=None, degrade_opt=None, device='cuda', resolution=None): if augment_opt is None: augment_opt = AUGMENT_OPT self.augment_opt = augment_opt if kernel_opt is None: kernel_opt = KERNEL_OPT self.kernel_opt = kernel_opt if degrade_opt is None: degrade_opt = DEGRADE_OPT self.degrade_opt = degrade_opt if resolution is not None: self.degrade_opt['gt_size'] = resolution self.device = device self.jpeger = DiffJPEG(differentiable=False).to(self.device) self.usm_sharpener = USMSharp().to(self.device) # blur settings for the first degradation self.blur_kernel_size = kernel_opt['blur_kernel_size'] self.kernel_list = kernel_opt['kernel_list'] self.kernel_prob = kernel_opt['kernel_prob'] # a list for each kernel probability self.blur_sigma = kernel_opt['blur_sigma'] self.betag_range = kernel_opt['betag_range'] # betag used in generalized Gaussian blur kernels self.betap_range = kernel_opt['betap_range'] # betap used in plateau blur kernels self.sinc_prob = kernel_opt['sinc_prob'] # the probability for sinc filters # blur settings for the second degradation self.blur_kernel_size2 = kernel_opt['blur_kernel_size2'] self.kernel_list2 = kernel_opt['kernel_list2'] self.kernel_prob2 = kernel_opt['kernel_prob2'] self.blur_sigma2 = kernel_opt['blur_sigma2'] self.betag_range2 = kernel_opt['betag_range2'] self.betap_range2 = kernel_opt['betap_range2'] self.sinc_prob2 = kernel_opt['sinc_prob2'] # a final sinc filter self.final_sinc_prob = kernel_opt['final_sinc_prob'] self.kernel_range = [2 * v + 1 for v in range(3, 11)] # kernel size ranges from 7 to 21 # TODO: kernel range is now hard-coded, should be in the configure file self.pulse_tensor = torch.zeros(21, 21).float() # convolving with pulse tensor brings no blurry effect self.pulse_tensor[10, 10] = 1 def get_kernel(self): # ------------------------ Generate kernels (used in the first degradation) ------------------------ # kernel_size = random.choice(self.kernel_range) if np.random.uniform() < self.kernel_opt['sinc_prob']: # this sinc filter setting is for kernels ranging from [7, 21] if kernel_size < 13: omega_c = np.random.uniform(np.pi / 3, np.pi) else: omega_c = np.random.uniform(np.pi / 5, np.pi) kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False) else: kernel = random_mixed_kernels( self.kernel_list, self.kernel_prob, kernel_size, self.blur_sigma, self.blur_sigma, [-math.pi, math.pi], self.betag_range, self.betap_range, noise_range=None) # pad kernel pad_size = (21 - kernel_size) // 2 kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size))) # ------------------------ Generate kernels (used in the second degradation) ------------------------ # kernel_size = random.choice(self.kernel_range) if np.random.uniform() < self.kernel_opt['sinc_prob2']: if kernel_size < 13: omega_c = np.random.uniform(np.pi / 3, np.pi) else: omega_c = np.random.uniform(np.pi / 5, np.pi) kernel2 = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False) else: kernel2 = random_mixed_kernels( self.kernel_list2, self.kernel_prob2, kernel_size, self.blur_sigma2, self.blur_sigma2, [-math.pi, math.pi], self.betag_range2, self.betap_range2, noise_range=None) # pad kernel pad_size = (21 - kernel_size) // 2 kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size))) # ------------------------------------- the final sinc kernel ------------------------------------- # if np.random.uniform() < self.kernel_opt['final_sinc_prob']: kernel_size = random.choice(self.kernel_range) omega_c = np.random.uniform(np.pi / 3, np.pi) sinc_kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=21) sinc_kernel = torch.FloatTensor(sinc_kernel) else: sinc_kernel = self.pulse_tensor # BGR to RGB, HWC to CHW, numpy to tensor kernel = torch.FloatTensor(kernel) kernel2 = torch.FloatTensor(kernel2) return (kernel, kernel2, sinc_kernel) @torch.no_grad() def __call__(self, img_gt, kernels=None): ''' :param: img_gt: BCHW, RGB, [0, 1] float32 tensor ''' if kernels is None: kernel = [] kernel2 = [] sinc_kernel = [] for _ in range(img_gt.shape[0]): k, k2, sk = self.get_kernel() kernel.append(k) kernel2.append(k2) sinc_kernel.append(sk) kernel = torch.stack(kernel) kernel2 = torch.stack(kernel2) sinc_kernel = torch.stack(sinc_kernel) else: # kernels created in dataset. kernel, kernel2, sinc_kernel = kernels # ----------------------- Pre-process ----------------------- # im_gt = img_gt.to(self.device) if self.degrade_opt['use_usm']: im_gt = self.usm_sharpener(im_gt) im_gt = im_gt.to(memory_format=torch.contiguous_format).float() kernel = kernel.to(self.device) kernel2 = kernel2.to(self.device) sinc_kernel = sinc_kernel.to(self.device) ori_h, ori_w = im_gt.size()[2:4] # ----------------------- The first degradation process ----------------------- # # blur out = filter2D(im_gt, kernel) # random resize updown_type = random.choices( ['up', 'down', 'keep'], self.degrade_opt['resize_prob'], )[0] if updown_type == 'up': scale = random.uniform(1, self.degrade_opt['resize_range'][1]) elif updown_type == 'down': scale = random.uniform(self.degrade_opt['resize_range'][0], 1) else: scale = 1 mode = random.choice(['area', 'bilinear', 'bicubic']) out = torch.nn.functional.interpolate(out, scale_factor=scale, mode=mode) # add noise gray_noise_prob = self.degrade_opt['gray_noise_prob'] if random.random() < self.degrade_opt['gaussian_noise_prob']: out = random_add_gaussian_noise_pt( out, sigma_range=self.degrade_opt['noise_range'], clip=True, rounds=False, gray_prob=gray_noise_prob, ) else: out = random_add_poisson_noise_pt( out, scale_range=self.degrade_opt['poisson_scale_range'], gray_prob=gray_noise_prob, clip=True, rounds=False) # JPEG compression jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.degrade_opt['jpeg_range']) out = torch.clamp(out, 0, 1) # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts out = self.jpeger(out, quality=jpeg_p) # ----------------------- The second degradation process ----------------------- # # blur if random.random() < self.degrade_opt['second_blur_prob']: out = out.contiguous() out = filter2D(out, kernel2) # random resize updown_type = random.choices( ['up', 'down', 'keep'], self.degrade_opt['resize_prob2'], )[0] if updown_type == 'up': scale = random.uniform(1, self.degrade_opt['resize_range2'][1]) elif updown_type == 'down': scale = random.uniform(self.degrade_opt['resize_range2'][0], 1) else: scale = 1 mode = random.choice(['area', 'bilinear', 'bicubic']) out = torch.nn.functional.interpolate( out, size=(int(ori_h / self.degrade_opt['sf'] * scale), int(ori_w / self.degrade_opt['sf'] * scale)), mode=mode, ) # add noise gray_noise_prob = self.degrade_opt['gray_noise_prob2'] if random.random() < self.degrade_opt['gaussian_noise_prob2']: out = random_add_gaussian_noise_pt( out, sigma_range=self.degrade_opt['noise_range2'], clip=True, rounds=False, gray_prob=gray_noise_prob, ) else: out = random_add_poisson_noise_pt( out, scale_range=self.degrade_opt['poisson_scale_range2'], gray_prob=gray_noise_prob, clip=True, rounds=False, ) # JPEG compression + the final sinc filter # We also need to resize images to desired sizes. We group [resize back + sinc filter] together # as one operation. # We consider two orders: # 1. [resize back + sinc filter] + JPEG compression # 2. JPEG compression + [resize back + sinc filter] # Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines. if random.random() < 0.5: # resize back + the final sinc filter mode = random.choice(['area', 'bilinear', 'bicubic']) out = torch.nn.functional.interpolate( out, size=(ori_h // self.degrade_opt['sf'], ori_w // self.degrade_opt['sf']), mode=mode, ) out = out.contiguous() out = filter2D(out, sinc_kernel) # JPEG compression jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.degrade_opt['jpeg_range2']) out = torch.clamp(out, 0, 1) out = self.jpeger(out, quality=jpeg_p) else: # JPEG compression jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.degrade_opt['jpeg_range2']) out = torch.clamp(out, 0, 1) out = self.jpeger(out, quality=jpeg_p) # resize back + the final sinc filter mode = random.choice(['area', 'bilinear', 'bicubic']) out = torch.nn.functional.interpolate( out, size=(ori_h // self.degrade_opt['sf'], ori_w // self.degrade_opt['sf']), mode=mode, ) out = out.contiguous() out = filter2D(out, sinc_kernel) # clamp and round im_lq = torch.clamp(out, 0, 1.0) # random crop gt_size = self.degrade_opt['gt_size'] im_gt, im_lq = paired_random_crop(im_gt, im_lq, gt_size, self.degrade_opt['sf']) if self.degrade_opt['resize_lq']: im_lq = torch.nn.functional.interpolate( im_lq, size=(im_gt.size(-2), im_gt.size(-1)), mode='bicubic', ) if random.random() < self.degrade_opt['no_degradation_prob'] or torch.isnan(im_lq).any(): im_lq = im_gt # sharpen self.gt again, as we have changed the self.gt with self._dequeue_and_enqueue im_lq = im_lq.contiguous() # for the warning: grad and param do not obey the gradient layout contract im_lq = im_lq*2 - 1.0 im_gt = im_gt*2 - 1.0 if self.degrade_opt['random_size']: raise NotImplementedError im_lq, im_gt = self.randn_cropinput(im_lq, im_gt) im_lq = torch.clamp(im_lq, -1.0, 1.0) im_gt = torch.clamp(im_gt, -1.0, 1.0) return (im_lq, im_gt)