mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add support for multiple LQ paths

Browse Source 
I want to be able to specify many different transformations onto
the target data; the model should handle them all. Do this by
allowing multiple LQ paths to be selected and the dataset class
selects one at random.
This commit is contained in:
James Betker 2020-05-06 17:24:17 -06:00
parent 3cd85f8073
commit 5c1832e124
1 changed files with 17 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
codes/data/LQGT_dataset.py
View File

@ -13,6 +13,10 @@ class LQGTDataset(data.Dataset):
If only GT images are provided, generate LQ images on-the-fly. If only GT images are provided, generate LQ images on-the-fly.
""" """
def get_lq_path(self, i):
which_lq = random.randint(0, len(self.paths_LQ))
return self.paths_LQ[which_lq][i]
def __init__(self, opt): def __init__(self, opt):
super(LQGTDataset, self).__init__() super(LQGTDataset, self).__init__()
self.opt = opt self.opt = opt
@ -23,17 +27,26 @@ class LQGTDataset(data.Dataset):
self.LQ_env, self.GT_env, self.PIX_env = None, None, None # environments for lmdbs self.LQ_env, self.GT_env, self.PIX_env = None, None, None # environments for lmdbs
self.paths_GT, self.sizes_GT = util.get_image_paths(self.data_type, opt['dataroot_GT']) self.paths_GT, self.sizes_GT = util.get_image_paths(self.data_type, opt['dataroot_GT'])
self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ']) self.paths_LQ = []
if isinstance(opt['dataroot_LQ'], list):
# Multiple LQ data sources can be given, in case there are multiple ways of corrupting a source image and
# we want the model to learn them all.
for dr_lq in opt['dataroot_LQ']:
lq_path, self.sizes_LQ = util.get_image_paths(self.data_type, dr_lq)
self.paths_LQ.append(lq_path)
else:
lq_path, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ'])
self.paths_LQ.append(lq_path)
self.doCrop = opt['doCrop'] self.doCrop = opt['doCrop']
if 'dataroot_PIX' in opt.keys(): if 'dataroot_PIX' in opt.keys():
self.paths_PIX, self.sizes_PIX = util.get_image_paths(self.data_type, opt['dataroot_PIX']) self.paths_PIX, self.sizes_PIX = util.get_image_paths(self.data_type, opt['dataroot_PIX'])
assert self.paths_GT, 'Error: GT path is empty.' assert self.paths_GT, 'Error: GT path is empty.'
if self.paths_LQ and self.paths_GT: if self.paths_LQ and self.paths_GT:
assert len(self.paths_LQ) == len( assert len(self.paths_LQ[0]) == len(
self.paths_GT self.paths_GT
), 'GT and LQ datasets have different number of images - {}, {}.'.format( ), 'GT and LQ datasets have different number of images - {}, {}.'.format(
len(self.paths_LQ), len(self.paths_GT)) len(self.paths_LQ[0]), len(self.paths_GT))
self.random_scale_list = [1] self.random_scale_list = [1]
def _init_lmdb(self): def _init_lmdb(self):
@ -74,7 +87,7 @@ class LQGTDataset(data.Dataset):
# get LQ image # get LQ image
if self.paths_LQ: if self.paths_LQ:
LQ_path = self.paths_LQ[index] LQ_path = self.get_lq_path(index)
resolution = [int(s) for s in self.sizes_LQ[index].split('_') resolution = [int(s) for s in self.sizes_LQ[index].split('_')
] if self.data_type == 'lmdb' else None ] if self.data_type == 'lmdb' else None
img_LQ = util.read_img(self.LQ_env, LQ_path, resolution) img_LQ = util.read_img(self.LQ_env, LQ_path, resolution)