|
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"] |
|
|
|
|
|
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: |
|
|
|
|
|
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] |
|
|
|
|
|
self.blur_kernel_size = opt["blur_kernel_size"] |
|
self.kernel_list = opt["kernel_list"] |
|
self.kernel_prob = opt["kernel_prob"] |
|
self.blur_sigma = opt["blur_sigma"] |
|
self.betag_range = opt[ |
|
"betag_range" |
|
] |
|
self.betap_range = opt["betap_range"] |
|
self.sinc_prob = opt["sinc_prob"] |
|
|
|
|
|
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"] |
|
|
|
|
|
self.final_sinc_prob = opt["final_sinc_prob"] |
|
|
|
self.kernel_range = [ |
|
2 * v + 1 for v in range(3, 11) |
|
] |
|
|
|
self.pulse_tensor = torch.zeros( |
|
21, 21 |
|
).float() |
|
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 |
|
) |
|
|
|
|
|
|
|
gt_path = self.paths[index] |
|
|
|
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}" |
|
) |
|
|
|
index = random.randint(0, self.__len__()) |
|
gt_path = self.paths[index] |
|
time.sleep(1) |
|
else: |
|
break |
|
finally: |
|
retry -= 1 |
|
img_gt = imfrombytes(img_bytes, float32=True) |
|
|
|
|
|
img_gt = augment(img_gt, self.opt["use_hflip"], self.opt["use_rot"]) |
|
|
|
|
|
|
|
h, w = img_gt.shape[0:2] |
|
crop_pad_size = 400 |
|
|
|
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 |
|
) |
|
|
|
if img_gt.shape[0] > crop_pad_size or img_gt.shape[1] > crop_pad_size: |
|
h, w = img_gt.shape[0:2] |
|
|
|
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, ...] |
|
|
|
|
|
kernel_size = random.choice(self.kernel_range) |
|
if np.random.uniform() < self.opt["sinc_prob"]: |
|
|
|
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_size = (21 - kernel_size) // 2 |
|
kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size))) |
|
|
|
|
|
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_size = (21 - kernel_size) // 2 |
|
kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size))) |
|
|
|
|
|
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 |
|
|
|
|
|
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) |
|
|