|
|
from importlib import import_module
|
|
|
|
|
|
from torch.utils.data import dataloader
|
|
|
from torch.utils.data import ConcatDataset
|
|
|
|
|
|
|
|
|
class MyConcatDataset(ConcatDataset):
|
|
|
def __init__(self, datasets):
|
|
|
super(MyConcatDataset, self).__init__(datasets)
|
|
|
self.train = datasets[0].train
|
|
|
|
|
|
def set_scale(self, idx_scale):
|
|
|
for d in self.datasets:
|
|
|
if hasattr(d, 'set_scale'): d.set_scale(idx_scale)
|
|
|
|
|
|
class Data:
|
|
|
def __init__(self, args):
|
|
|
self.loader_train = None
|
|
|
if not args.test_only:
|
|
|
datasets = []
|
|
|
for d in args.data_train:
|
|
|
module_name = d if d.find('DIV2K-Q') < 0 else 'DIV2KJPEG'
|
|
|
m = import_module('data.' + module_name.lower())
|
|
|
datasets.append(getattr(m, module_name)(args, name=d))
|
|
|
|
|
|
self.loader_train = dataloader.DataLoader(
|
|
|
MyConcatDataset(datasets),
|
|
|
batch_size=args.batch_size,
|
|
|
shuffle=True,
|
|
|
pin_memory=not args.cpu,
|
|
|
num_workers=args.n_threads,
|
|
|
)
|
|
|
|
|
|
self.loader_test = []
|
|
|
for d in args.data_test:
|
|
|
if d in ['CBSD68','Kodak24','McM','Set5', 'Set14', 'B100', 'Urban100']:
|
|
|
m = import_module('data.benchmark')
|
|
|
testset = getattr(m, 'Benchmark')(args, train=False, name=d)
|
|
|
else:
|
|
|
module_name = d if d.find('DIV2K-Q') < 0 else 'DIV2KJPEG'
|
|
|
m = import_module('data.' + module_name.lower())
|
|
|
testset = getattr(m, module_name)(args, train=False, name=d)
|
|
|
|
|
|
self.loader_test.append(
|
|
|
dataloader.DataLoader(
|
|
|
testset,
|
|
|
batch_size=1,
|
|
|
shuffle=False,
|
|
|
pin_memory=not args.cpu,
|
|
|
num_workers=args.n_threads,
|
|
|
)
|
|
|
)
|
|
|
|
