import cv2 import math import numpy as np import os import os.path as osp import random import time import torch from basicsr.data.degradations import circular_lowpass_kernel, random_mixed_kernels from basicsr.data.transforms import augment from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor from basicsr.utils.registry import DATASET_REGISTRY from torch.utils import data as data @DATASET_REGISTRY.register() class RealESRGANDataset(data.Dataset): """Dataset used for Real-ESRGAN model: Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data. It loads gt (Ground-Truth) images, and augments them. It also generates blur kernels and sinc kernels for generating low-quality images. Note that the low-quality images are processed in tensors on GPUS for faster processing. Args: opt (dict): Config for train datasets. It contains the following keys: dataroot_gt (str): Data root path for gt. meta_info (str): Path for meta information file. io_backend (dict): IO backend type and other kwarg. use_hflip (bool): Use horizontal flips. use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation). Please see more options in the codes. """ def __init__(self, opt): super(RealESRGANDataset, self).__init__() self.opt = opt self.file_client = None self.io_backend_opt = opt["io_backend"] self.gt_folder = opt["dataroot_gt"] # file client (lmdb io backend) if self.io_backend_opt["type"] == "lmdb": self.io_backend_opt["db_paths"] = [self.gt_folder] self.io_backend_opt["client_keys"] = ["gt"] if not self.gt_folder.endswith(".lmdb"): raise ValueError( f"'dataroot_gt' should end with '.lmdb', but received {self.gt_folder}" ) with open(osp.join(self.gt_folder, "meta_info.txt")) as fin: self.paths = [line.split(".")[0] for line in fin] else: # disk backend with meta_info # Each line in the meta_info describes the relative path to an image with open(self.opt["meta_info"]) as fin: paths = [line.strip().split(" ")[0] for line in fin] self.paths = [os.path.join(self.gt_folder, v) for v in paths] # blur settings for the first degradation self.blur_kernel_size = opt["blur_kernel_size"] self.kernel_list = opt["kernel_list"] self.kernel_prob = opt["kernel_prob"] # a list for each kernel probability self.blur_sigma = opt["blur_sigma"] self.betag_range = opt[ "betag_range" ] # betag used in generalized Gaussian blur kernels self.betap_range = opt["betap_range"] # betap used in plateau blur kernels self.sinc_prob = opt["sinc_prob"] # the probability for sinc filters # blur settings for the second degradation self.blur_kernel_size2 = opt["blur_kernel_size2"] self.kernel_list2 = opt["kernel_list2"] self.kernel_prob2 = opt["kernel_prob2"] self.blur_sigma2 = opt["blur_sigma2"] self.betag_range2 = opt["betag_range2"] self.betap_range2 = opt["betap_range2"] self.sinc_prob2 = opt["sinc_prob2"] # a final sinc filter self.final_sinc_prob = 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 __getitem__(self, index): if self.file_client is None: self.file_client = FileClient( self.io_backend_opt.pop("type"), **self.io_backend_opt ) # -------------------------------- Load gt images -------------------------------- # # Shape: (h, w, c); channel order: BGR; image range: [0, 1], float32. gt_path = self.paths[index] # avoid errors caused by high latency in reading files retry = 3 while retry > 0: try: img_bytes = self.file_client.get(gt_path, "gt") except (IOError, OSError) as e: logger = get_root_logger() logger.warn( f"File client error: {e}, remaining retry times: {retry - 1}" ) # change another file to read index = random.randint(0, self.__len__()) gt_path = self.paths[index] time.sleep(1) # sleep 1s for occasional server congestion else: break finally: retry -= 1 img_gt = imfrombytes(img_bytes, float32=True) # -------------------- Do augmentation for training: flip, rotation -------------------- # img_gt = augment(img_gt, self.opt["use_hflip"], self.opt["use_rot"]) # crop or pad to 400 # TODO: 400 is hard-coded. You may change it accordingly h, w = img_gt.shape[0:2] crop_pad_size = 400 # pad if h < crop_pad_size or w < crop_pad_size: pad_h = max(0, crop_pad_size - h) pad_w = max(0, crop_pad_size - w) img_gt = cv2.copyMakeBorder( img_gt, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT_101 ) # crop if img_gt.shape[0] > crop_pad_size or img_gt.shape[1] > crop_pad_size: h, w = img_gt.shape[0:2] # randomly choose top and left coordinates top = random.randint(0, h - crop_pad_size) left = random.randint(0, w - crop_pad_size) img_gt = img_gt[top : top + crop_pad_size, left : left + crop_pad_size, ...] # ------------------------ Generate kernels (used in the first degradation) ------------------------ # kernel_size = random.choice(self.kernel_range) if np.random.uniform() < self.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.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.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 img_gt = img2tensor([img_gt], bgr2rgb=True, float32=True)[0] kernel = torch.FloatTensor(kernel) kernel2 = torch.FloatTensor(kernel2) return_d = { "gt": img_gt, "kernel1": kernel, "kernel2": kernel2, "sinc_kernel": sinc_kernel, "gt_path": gt_path, } return return_d def __len__(self): return len(self.paths)