Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
10 changed files
with
1,545 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,70 @@ | ||
from torch import nn as nn | ||
from torch.nn import functional as F | ||
|
||
from basicsr.utils.registry import ARCH_REGISTRY | ||
|
||
|
||
@ARCH_REGISTRY.register(suffix='basicsr') | ||
class SRVGGNetCompact(nn.Module): | ||
"""A compact VGG-style network structure for super-resolution. | ||
It is a compact network structure, which performs upsampling in the last layer and no convolution is | ||
conducted on the HR feature space. | ||
Args: | ||
num_in_ch (int): Channel number of inputs. Default: 3. | ||
num_out_ch (int): Channel number of outputs. Default: 3. | ||
num_feat (int): Channel number of intermediate features. Default: 64. | ||
num_conv (int): Number of convolution layers in the body network. Default: 16. | ||
upscale (int): Upsampling factor. Default: 4. | ||
act_type (str): Activation type, options: 'relu', 'prelu', 'leakyrelu'. Default: prelu. | ||
""" | ||
|
||
def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu'): | ||
super(SRVGGNetCompact, self).__init__() | ||
self.num_in_ch = num_in_ch | ||
self.num_out_ch = num_out_ch | ||
self.num_feat = num_feat | ||
self.num_conv = num_conv | ||
self.upscale = upscale | ||
self.act_type = act_type | ||
|
||
self.body = nn.ModuleList() | ||
# the first conv | ||
self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)) | ||
# the first activation | ||
if act_type == 'relu': | ||
activation = nn.ReLU(inplace=True) | ||
elif act_type == 'prelu': | ||
activation = nn.PReLU(num_parameters=num_feat) | ||
elif act_type == 'leakyrelu': | ||
activation = nn.LeakyReLU(negative_slope=0.1, inplace=True) | ||
self.body.append(activation) | ||
|
||
# the body structure | ||
for _ in range(num_conv): | ||
self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1)) | ||
# activation | ||
if act_type == 'relu': | ||
activation = nn.ReLU(inplace=True) | ||
elif act_type == 'prelu': | ||
activation = nn.PReLU(num_parameters=num_feat) | ||
elif act_type == 'leakyrelu': | ||
activation = nn.LeakyReLU(negative_slope=0.1, inplace=True) | ||
self.body.append(activation) | ||
|
||
# the last conv | ||
self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1)) | ||
# upsample | ||
self.upsampler = nn.PixelShuffle(upscale) | ||
|
||
def forward(self, x): | ||
out = x | ||
for i in range(0, len(self.body)): | ||
out = self.body[i](out) | ||
|
||
out = self.upsampler(out) | ||
# add the nearest upsampled image, so that the network learns the residual | ||
base = F.interpolate(x, scale_factor=self.upscale, mode='nearest') | ||
out += base | ||
return out |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,193 @@ | ||
import cv2 | ||
import math | ||
import numpy as np | ||
import os | ||
import os.path as osp | ||
import random | ||
import time | ||
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 FileClient, get_root_logger, imfrombytes, img2tensor | ||
from basicsr.utils.registry import DATASET_REGISTRY | ||
|
||
|
||
@DATASET_REGISTRY.register(suffix='basicsr') | ||
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) |
Oops, something went wrong.