|
import os
|
|
import random
|
|
import imageio
|
|
import numpy as np
|
|
import torch.utils.data as data
|
|
|
|
from data import common
|
|
|
|
from utils import interact
|
|
|
|
class Dataset(data.Dataset):
|
|
"""Basic dataloader class
|
|
"""
|
|
def __init__(self, args, mode='train'):
|
|
super(Dataset, self).__init__()
|
|
self.args = args
|
|
self.mode = mode
|
|
|
|
self.modes = ()
|
|
self.set_modes()
|
|
self._check_mode()
|
|
|
|
self.set_keys()
|
|
|
|
if self.mode == 'train':
|
|
dataset = args.data_train
|
|
elif self.mode == 'val':
|
|
dataset = args.data_val
|
|
elif self.mode == 'test':
|
|
dataset = args.data_test
|
|
elif self.mode == 'demo':
|
|
pass
|
|
else:
|
|
raise NotImplementedError('not implemented for this mode: {}!'.format(self.mode))
|
|
|
|
if self.mode == 'demo':
|
|
self.subset_root = args.demo_input_dir
|
|
else:
|
|
self.subset_root = os.path.join(args.data_root, dataset, self.mode)
|
|
|
|
self.blur_list = []
|
|
self.sharp_list = []
|
|
|
|
self._scan()
|
|
|
|
def set_modes(self):
|
|
self.modes = ('train', 'val', 'test', 'demo')
|
|
|
|
def _check_mode(self):
|
|
"""Should be called in the child class __init__() after super
|
|
"""
|
|
if self.mode not in self.modes:
|
|
raise NotImplementedError('mode error: not for {}'.format(self.mode))
|
|
|
|
return
|
|
|
|
def set_keys(self):
|
|
self.blur_key = 'blur'
|
|
self.sharp_key = 'sharp'
|
|
|
|
self.non_blur_keys = []
|
|
self.non_sharp_keys = []
|
|
|
|
return
|
|
|
|
def _scan(self, root=None):
|
|
"""Should be called in the child class __init__() after super
|
|
"""
|
|
if root is None:
|
|
root = self.subset_root
|
|
|
|
if self.blur_key in self.non_blur_keys:
|
|
self.non_blur_keys.remove(self.blur_key)
|
|
if self.sharp_key in self.non_sharp_keys:
|
|
self.non_sharp_keys.remove(self.sharp_key)
|
|
|
|
def _key_check(path, true_key, false_keys):
|
|
path = os.path.join(path, '')
|
|
if path.find(true_key) >= 0:
|
|
for false_key in false_keys:
|
|
if path.find(false_key) >= 0:
|
|
return False
|
|
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def _get_list_by_key(root, true_key, false_keys):
|
|
data_list = []
|
|
for sub, dirs, files in os.walk(root):
|
|
if not dirs:
|
|
file_list = [os.path.join(sub, f) for f in files]
|
|
if _key_check(sub, true_key, false_keys):
|
|
data_list += file_list
|
|
|
|
data_list.sort()
|
|
|
|
return data_list
|
|
|
|
def _rectify_keys():
|
|
self.blur_key = os.path.join(self.blur_key, '')
|
|
self.non_blur_keys = [os.path.join(non_blur_key, '') for non_blur_key in self.non_blur_keys]
|
|
self.sharp_key = os.path.join(self.sharp_key, '')
|
|
self.non_sharp_keys = [os.path.join(non_sharp_key, '') for non_sharp_key in self.non_sharp_keys]
|
|
|
|
_rectify_keys()
|
|
|
|
self.blur_list = _get_list_by_key(root, self.blur_key, self.non_blur_keys)
|
|
self.sharp_list = _get_list_by_key(root, self.sharp_key, self.non_sharp_keys)
|
|
|
|
if len(self.sharp_list) > 0:
|
|
assert(len(self.blur_list) == len(self.sharp_list))
|
|
|
|
return
|
|
|
|
def __getitem__(self, idx):
|
|
|
|
blur = imageio.imread(self.blur_list[idx], pilmode='RGB')
|
|
if len(self.sharp_list) > 0:
|
|
sharp = imageio.imread(self.sharp_list[idx], pilmode='RGB')
|
|
imgs = [blur, sharp]
|
|
else:
|
|
imgs = [blur]
|
|
|
|
pad_width = 0
|
|
if self.mode == 'train':
|
|
imgs = common.crop(*imgs, ps=self.args.patch_size)
|
|
if self.args.augment:
|
|
imgs = common.augment(*imgs, hflip=True, rot=True, shuffle=True, change_saturation=True, rgb_range=self.args.rgb_range)
|
|
imgs[0] = common.add_noise(imgs[0], sigma_sigma=2, rgb_range=self.args.rgb_range)
|
|
elif self.mode == 'demo':
|
|
imgs[0], pad_width = common.pad(imgs[0], divisor=2**(self.args.n_scales-1))
|
|
else:
|
|
pass
|
|
|
|
if self.args.gaussian_pyramid:
|
|
imgs = common.generate_pyramid(*imgs, n_scales=self.args.n_scales)
|
|
|
|
imgs = common.np2tensor(*imgs, rgb_range=self.args.rgb_range)
|
|
relpath = os.path.relpath(self.blur_list[idx], self.subset_root)
|
|
|
|
blur = imgs[0]
|
|
sharp = imgs[1] if len(imgs) > 1 else False
|
|
|
|
return blur, sharp, pad_width, idx, relpath
|
|
|
|
def __len__(self):
|
|
return len(self.blur_list)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
