Spaces:
Build error
Build error
| import math | |
| from random import choice, choices, randint | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| from torch.utils.data.dataset import Dataset | |
| from utils import USMSharp_npy, cvtColor, preprocess_input | |
| from .degradations import (circular_lowpass_kernel, random_add_gaussian_noise, | |
| random_add_poisson_noise, random_mixed_kernels) | |
| from .transforms import augment, paired_random_crop | |
| def cv_show(image): | |
| image = np.array(image) | |
| image = cv2.resize(image, (256, 128), interpolation=cv2.INTER_CUBIC) | |
| cv2.imshow('image', image) | |
| cv2.waitKey(0) | |
| cv2.destroyAllWindows() | |
| def get_new_img_size(width, height, img_min_side=600): | |
| if width <= height: | |
| f = float(img_min_side) / width | |
| resized_height = int(f * height) | |
| resized_width = int(img_min_side) | |
| else: | |
| f = float(img_min_side) / height | |
| resized_width = int(f * width) | |
| resized_height = int(img_min_side) | |
| return resized_width, resized_height | |
| class SRGANDataset(Dataset): | |
| def __init__(self, train_lines, lr_shape, hr_shape): | |
| super(SRGANDataset, self).__init__() | |
| self.train_lines = train_lines | |
| self.train_batches = len(train_lines) | |
| self.lr_shape = lr_shape | |
| self.hr_shape = hr_shape | |
| self.scale = int(hr_shape[0]/lr_shape[0]) | |
| self.usmsharp = USMSharp_npy() | |
| # 第一次滤波的参数 | |
| self.blur_kernel_size = 21 | |
| self.kernel_list = ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso'] | |
| self.kernel_prob = [0.45, 0.25, 0.12, 0.03, 0.12, 0.03] | |
| self.sinc_prob = 0.1 | |
| self.blur_sigma = [0.2, 3] | |
| self.betag_range = [0.5, 4] | |
| self.betap_range = [1, 2] | |
| # 第二次滤波的参数 | |
| self.blur_kernel_size2 = 21 | |
| self.kernel_list2 = ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso'] | |
| self.kernel_prob2 = [0.45, 0.25, 0.12, 0.03, 0.12, 0.03] | |
| self.sinc_prob2 = 0.1 | |
| self.blur_sigma2 = [0.2, 3] | |
| self.betag_range2 = [0.5, 4] | |
| self.betap_range2 = [1, 2] | |
| # 最后的sinc滤波 | |
| self.final_sinc_prob = 0.8 | |
| # 卷积核大小从7到21分布 | |
| self.kernel_range = [2 * v + 1 for v in range(3, 11)] | |
| # 使用脉冲张量进行卷积不会产生模糊效果 | |
| self.pulse_tensor = np.zeros(shape=[21, 21], dtype='float32') | |
| self.pulse_tensor[10, 10] = 1 | |
| # 第一次退化的参数 | |
| self.resize_prob = [0.2, 0.7, 0.1] # up, down, keep | |
| self.resize_range = [0.15, 1.5] | |
| self.gaussian_noise_prob = 0.5 | |
| self.noise_range = [1, 30] | |
| self.poisson_scale_range = [0.05, 3] | |
| self.gray_noise_prob = 0.4 | |
| self.jpeg_range = [30, 95] | |
| # 第二次退化的参数 | |
| self.second_blur_prob = 0.8 | |
| self.resize_prob2 = [0.3, 0.4, 0.3] # up, down, keep | |
| self.resize_range2 = [0.3, 1.2] | |
| self.gaussian_noise_prob2= 0.5 | |
| self.noise_range2 = [1, 25] | |
| self.poisson_scale_range2= [0.05, 2.5] | |
| self.gray_noise_prob2 = 0.4 | |
| self.jpeg_range2 = [30, 95] | |
| def __len__(self): | |
| return self.train_batches | |
| def __getitem__(self, index): | |
| index = index % self.train_batches | |
| image_origin = Image.open(self.train_lines[index].split()[0]) | |
| lq, gt = self.get_random_data(image_origin, self.hr_shape) | |
| gt = np.transpose(preprocess_input(np.array(gt, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1]) | |
| lq = np.transpose(preprocess_input(np.array(lq, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1]) | |
| return lq, gt | |
| def rand(self, a=0, b=1): | |
| return np.random.rand()*(b-a) + a | |
| def get_random_data(self, image, input_shape): | |
| #------------------------------# | |
| # 读取图像并转换成RGB图像 | |
| # cvtColor将np转Image | |
| #------------------------------# | |
| image = cvtColor(image) | |
| #------------------------------# | |
| # 获得图像的高宽与目标高宽 | |
| #------------------------------# | |
| iw, ih = image.size | |
| h, w = input_shape | |
| scale = min(w/iw, h/ih) | |
| nw = int(iw*scale) | |
| nh = int(ih*scale) | |
| dx = (w-nw)//2 | |
| dy = (h-nh)//2 | |
| #---------------------------------# | |
| # 将图像多余的部分加上灰条 | |
| #---------------------------------# | |
| image = image.resize((nw,nh), Image.BICUBIC) | |
| new_image = Image.new('RGB', (w,h), (128,128,128)) | |
| new_image.paste(image, (dx, dy)) | |
| image = np.array(new_image, np.float32) | |
| rotate = self.rand()<.5 | |
| if rotate: | |
| angle = np.random.randint(-15,15) | |
| a,b = w/2,h/2 | |
| M = cv2.getRotationMatrix2D((a,b),angle,1) | |
| image = cv2.warpAffine(np.array(image), M, (w,h), borderValue=[128,128,128]) | |
| # ------------------------ 生成卷积核以进行第一次退化处理 ------------------------ # | |
| kernel_size = choice(self.kernel_range) | |
| if np.random.uniform() < self.sinc_prob: | |
| # 此sinc过滤器设置适用于[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))) | |
| kernel = kernel.astype(np.float32) | |
| # ------------------------ 生成卷积核以进行第二次退化处理 ------------------------ # | |
| kernel_size = choice(self.kernel_range) | |
| if np.random.uniform() < self.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))) | |
| kernel2 = kernel2.astype(np.float32) | |
| # ----------------------the final sinc kernel ------------------------- # | |
| if np.random.uniform() < self.final_sinc_prob: | |
| kernel_size = 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) | |
| else: | |
| sinc_kernel = self.pulse_tensor | |
| sinc_kernel = sinc_kernel.astype(np.float32) | |
| lq, gt = self.feed_data(image, kernel, kernel2, sinc_kernel) | |
| return lq, gt | |
| def feed_data(self, img_gt, kernel1, kernel2, sinc_kernel): | |
| img_gt = np.array(img_gt, dtype=np.float32) | |
| # 对gt进行锐化 | |
| img_gt = np.clip(img_gt / 255, 0, 1) | |
| gt = self.usmsharp.filt(img_gt) | |
| [ori_w, ori_h, _] = gt.shape | |
| # ---------------------- 根据参数进行第一次退化 -------------------- # | |
| # 模糊处理 | |
| out = cv2.filter2D(img_gt, -1, kernel1) | |
| # 随机 resize | |
| updown_type = choices(['up', 'down', 'keep'], self.resize_prob)[0] | |
| if updown_type == 'up': | |
| scale = np.random.uniform(1, self.resize_range[1]) | |
| elif updown_type == 'down': | |
| scale = np.random.uniform(self.resize_range[0], 1) | |
| else: | |
| scale = 1 | |
| mode = choice(['area', 'bilinear', 'bicubic']) | |
| if mode=='area': | |
| out = cv2.resize(out, (int(ori_h * scale), int(ori_w * scale)), interpolation=cv2.INTER_AREA) | |
| elif mode=='bilinear': | |
| out = cv2.resize(out, (int(ori_h * scale), int(ori_w * scale)), interpolation=cv2.INTER_LINEAR) | |
| else: | |
| out = cv2.resize(out, (int(ori_h * scale), int(ori_w * scale)), interpolation=cv2.INTER_CUBIC) | |
| # 灰度噪声 | |
| gray_noise_prob = self.gray_noise_prob | |
| if np.random.uniform() < self.gaussian_noise_prob: | |
| out = random_add_gaussian_noise( | |
| out, sigma_range=self.noise_range, clip=True, rounds=False, gray_prob=gray_noise_prob) | |
| else: | |
| out = random_add_poisson_noise( | |
| out, | |
| scale_range=self.poisson_scale_range, | |
| gray_prob=gray_noise_prob, | |
| clip=True, | |
| rounds=False) | |
| # JPEG 压缩 | |
| jpeg_p = np.random.uniform(low=self.jpeg_range[0], high=self.jpeg_range[1]) | |
| jpeg_p = int(jpeg_p) | |
| out = np.clip(out, 0, 1) | |
| encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_p] | |
| _, encimg = cv2.imencode('.jpg', out * 255., encode_param) | |
| out = np.float32(cv2.imdecode(encimg, 1))/255 | |
| # ---------------------- 根据参数进行第一次退化 -------------------- # | |
| # 模糊 | |
| if np.random.uniform() < self.second_blur_prob: | |
| out = cv2.filter2D(out, -1, kernel2) | |
| # 随机 resize | |
| updown_type = choices(['up', 'down', 'keep'], self.resize_prob2)[0] | |
| if updown_type == 'up': | |
| scale = np.random.uniform(1, self.resize_range2[1]) | |
| elif updown_type == 'down': | |
| scale = np.random.uniform(self.resize_range2[0], 1) | |
| else: | |
| scale = 1 | |
| mode = choice(['area', 'bilinear', 'bicubic']) | |
| if mode == 'area': | |
| out = cv2.resize(out, (int(ori_h / self.scale * scale), int(ori_w / self.scale * scale)), interpolation=cv2.INTER_AREA) | |
| elif mode == 'bilinear': | |
| out = cv2.resize(out, (int(ori_h / self.scale * scale), int(ori_w / self.scale * scale)), interpolation=cv2.INTER_LINEAR) | |
| else: | |
| out = cv2.resize(out, (int(ori_h / self.scale * scale), int(ori_w / self.scale * scale)), interpolation=cv2.INTER_CUBIC) | |
| # 灰度噪声 | |
| gray_noise_prob = self.gray_noise_prob2 | |
| if np.random.uniform() < self.gaussian_noise_prob2: | |
| out = random_add_gaussian_noise( | |
| out, sigma_range=self.noise_range2, clip=True, rounds=False, gray_prob=gray_noise_prob) | |
| else: | |
| out = random_add_poisson_noise( | |
| out, | |
| scale_range=self.poisson_scale_range2, | |
| gray_prob=gray_noise_prob, | |
| clip=True, | |
| rounds=False) | |
| # JPEG压缩+最后的sinc滤波器 | |
| # 我们还需要将图像的大小调整到所需的尺寸。我们把[调整大小+sinc过滤器]组合在一起 | |
| # 作为一个操作。 | |
| # 我们考虑两个顺序。 | |
| # 1. [调整大小+sinc filter] + JPEG压缩 | |
| # 2. 2. JPEG压缩+[调整大小+sinc过滤]。 | |
| # 根据经验，我们发现其他组合（sinc + JPEG + Resize）会引入扭曲的线条。 | |
| if np.random.uniform() < 0.5: | |
| # resize back + the final sinc filter | |
| mode = choice(['area', 'bilinear', 'bicubic']) | |
| if mode == 'area': | |
| out = cv2.resize(out, (ori_h // self.scale, ori_w // self.scale), interpolation=cv2.INTER_AREA) | |
| elif mode == 'bilinear': | |
| out = cv2.resize(out, (ori_h // self.scale, ori_w // self.scale), interpolation=cv2.INTER_LINEAR) | |
| else: | |
| out = cv2.resize(out, (ori_h // self.scale, ori_w // self.scale), interpolation=cv2.INTER_CUBIC) | |
| out = cv2.filter2D(out, -1, sinc_kernel) | |
| # JPEG 压缩 | |
| jpeg_p = np.random.uniform(low=self.jpeg_range[0], high=self.jpeg_range[1]) | |
| jpeg_p = jpeg_p | |
| out = np.clip(out, 0, 1) | |
| encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_p] | |
| _, encimg = cv2.imencode('.jpg', out * 255., encode_param) | |
| out = np.float32(cv2.imdecode(encimg, 1)) / 255 | |
| else: | |
| # JPEG 压缩 | |
| jpeg_p = np.random.uniform(low=self.jpeg_range[0], high=self.jpeg_range[1]) | |
| jpeg_p = jpeg_p | |
| out = np.clip(out, 0, 1) | |
| encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_p] | |
| _, encimg = cv2.imencode('.jpg', out * 255., encode_param) | |
| out = np.float32(cv2.imdecode(encimg, 1)) / 255 | |
| # resize back + the final sinc filter | |
| mode = choice(['area', 'bilinear', 'bicubic']) | |
| if mode == 'area': | |
| out = cv2.resize(out, (ori_h // self.scale, ori_w // self.scale),interpolation=cv2.INTER_AREA) | |
| elif mode == 'bilinear': | |
| out = cv2.resize(out, (ori_h // self.scale, ori_w // self.scale),interpolation=cv2.INTER_LINEAR) | |
| else: | |
| out = cv2.resize(out, (ori_h // self.scale, ori_w // self.scale),interpolation=cv2.INTER_CUBIC) | |
| lq = np.clip((out * 255.0), 0, 255) | |
| gt = np.clip((gt * 255.0), 0, 255) | |
| return Image.fromarray(np.uint8(lq)), Image.fromarray(np.uint8(gt)) | |
| def SRGAN_dataset_collate(batch): | |
| images_l = [] | |
| images_h = [] | |
| for img_l, img_h in batch: | |
| images_l.append(img_l) | |
| images_h.append(img_h) | |
| return np.array(images_l), np.array(images_h) | |