def main(mode='folder'):
"""Test vimeo90k dataset.
Args:
mode: There are two modes: 'lmdb', 'folder'.
"""
opt = {}
opt['dist'] = False
opt['phase'] = 'train'
opt['name'] = 'Vimeo90K'
opt['type'] = 'Vimeo90KDataset'
if mode == 'folder':
opt['dataroot_gt'] = 'datasets/vimeo90k/vimeo_septuplet/sequences'
opt['dataroot_lq'] = 'datasets/vimeo90k/vimeo_septuplet_matlabLRx4/sequences' # noqa E501
opt['meta_info_file'] = 'basicsr/data/meta_info/meta_info_Vimeo90K_train_GT.txt' # noqa E501
opt['io_backend'] = dict(type='disk')
elif mode == 'lmdb':
opt['dataroot_gt'] = 'datasets/vimeo90k/vimeo90k_train_GT_only4th.lmdb'
opt['dataroot_lq'] = 'datasets/vimeo90k/vimeo90k_train_LR7frames.lmdb'
opt['meta_info_file'] = 'basicsr/data/meta_info/meta_info_Vimeo90K_train_GT.txt' # noqa E501
opt['io_backend'] = dict(type='lmdb')
opt['num_frame'] = 7
opt['gt_size'] = 256
opt['random_reverse'] = True
opt['use_hflip'] = True
opt['use_rot'] = True
opt['use_shuffle'] = True
opt['num_worker_per_gpu'] = 1
opt['batch_size_per_gpu'] = 16
opt['scale'] = 4
opt['dataset_enlarge_ratio'] = 1
os.makedirs('tmp', exist_ok=True)
dataset = build_dataset(opt)
data_loader = build_dataloader(dataset, opt, num_gpu=0, dist=opt['dist'], sampler=None)
nrow = int(math.sqrt(opt['batch_size_per_gpu']))
padding = 2 if opt['phase'] == 'train' else 0
print('start...')
for i, data in enumerate(data_loader):
if i > 5:
break
print(i)
lq = data['lq']
gt = data['gt']
key = data['key']
print(key)
for j in range(opt['num_frame']):
torchvision.utils.save_image(
lq[:, j, :, :, :], f'tmp/lq_{i:03d}_frame{j}.png', nrow=nrow, padding=padding, normalize=False)
torchvision.utils.save_image(gt, f'tmp/gt_{i:03d}.png', nrow=nrow, padding=padding, normalize=False)
def create_train_val_dataloader(opt, logger):
# create train and val dataloaders
train_loader, val_loaders = None, []
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
dataset_enlarge_ratio = dataset_opt.get('dataset_enlarge_ratio', 1)
train_set = build_dataset(dataset_opt)
train_sampler = EnlargedSampler(train_set, opt['world_size'],
opt['rank'], dataset_enlarge_ratio)
train_loader = build_dataloader(train_set,
dataset_opt,
num_gpu=opt['num_gpu'],
dist=opt['dist'],
sampler=train_sampler,
seed=opt['manual_seed'])
num_iter_per_epoch = math.ceil(
len(train_set) * dataset_enlarge_ratio /
(dataset_opt['batch_size_per_gpu'] * opt['world_size']))
total_iters = int(opt['train']['total_iter'])
total_epochs = math.ceil(total_iters / (num_iter_per_epoch))
logger.info(
'Training statistics:'
f'\n\tNumber of train images: {len(train_set)}'
f'\n\tDataset enlarge ratio: {dataset_enlarge_ratio}'
f'\n\tBatch size per gpu: {dataset_opt["batch_size_per_gpu"]}'
f'\n\tWorld size (gpu number): {opt["world_size"]}'
f'\n\tRequire iter number per epoch: {num_iter_per_epoch}'
f'\n\tTotal epochs: {total_epochs}; iters: {total_iters}.')
elif phase.split('_')[0] == 'val':
val_set = build_dataset(dataset_opt)
val_loader = build_dataloader(val_set,
dataset_opt,
num_gpu=opt['num_gpu'],
dist=opt['dist'],
sampler=None,
seed=opt['manual_seed'])
logger.info(
f'Number of val images/folders in {dataset_opt["name"]}: {len(val_set)}'
)
val_loaders.append(val_loader)
else:
raise ValueError(f'Dataset phase {phase} is not recognized.')
return train_loader, train_sampler, val_loaders, total_epochs, total_iters
def main():
"""Test FFHQ dataset."""
opt = {}
opt['dist'] = False
opt['gpu_ids'] = [0]
opt['phase'] = 'train'
opt['name'] = 'FFHQ'
opt['type'] = 'FFHQDataset'
opt['dataroot_gt'] = 'datasets/ffhq/ffhq_256.lmdb'
opt['io_backend'] = dict(type='lmdb')
opt['use_hflip'] = True
opt['mean'] = [0.5, 0.5, 0.5]
opt['std'] = [0.5, 0.5, 0.5]
opt['use_shuffle'] = True
opt['num_worker_per_gpu'] = 1
opt['batch_size_per_gpu'] = 4
opt['dataset_enlarge_ratio'] = 1
os.makedirs('tmp', exist_ok=True)
dataset = build_dataset(opt)
data_loader = build_dataloader(dataset,
opt,
num_gpu=0,
dist=opt['dist'],
sampler=None)
nrow = int(math.sqrt(opt['batch_size_per_gpu']))
padding = 2 if opt['phase'] == 'train' else 0
print('start...')
for i, data in enumerate(data_loader):
if i > 5:
break
print(i)
gt = data['gt']
print(torch.min(gt), torch.max(gt))
gt_path = data['gt_path']
print(gt_path)
torchvision.utils.save_image(gt,
f'tmp/gt_{i:03d}.png',
nrow=nrow,
padding=padding,
normalize=True,
range=(-1, 1))
def test_pipeline(root_path):
# parse options, set distributed setting, set ramdom seed
opt = parse_options(root_path, is_train=False)
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
# mkdir and initialize loggers
make_exp_dirs(opt)
log_file = osp.join(opt['path']['log'],
f"test_{opt['name']}_{get_time_str()}.log")
logger = get_root_logger(logger_name='basicsr',
log_level=logging.INFO,
log_file=log_file)
logger.info(get_env_info())
logger.info(dict2str(opt))
# create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt['datasets'].items()):
test_set = build_dataset(dataset_opt)
test_loader = build_dataloader(test_set,
dataset_opt,
num_gpu=opt['num_gpu'],
dist=opt['dist'],
sampler=None,
seed=opt['manual_seed'])
logger.info(
f"Number of test images in {dataset_opt['name']}: {len(test_set)}")
test_loaders.append(test_loader)
# create model
model = build_model(opt)
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt['name']
logger.info(f'Testing {test_set_name}...')
model.validation(test_loader,
current_iter=opt['name'],
tb_logger=None,
save_img=opt['val']['save_img'])
def main(mode='folder'):
"""Test reds dataset.
Args:
mode: There are two modes: 'lmdb', 'folder'.
"""
opt = {}
opt['dist'] = False
opt['phase'] = 'train'
opt['name'] = 'REDS'
opt['type'] = 'REDSDataset'
if mode == 'folder':
opt['dataroot_gt'] = 'datasets/REDS/train_sharp'
opt['dataroot_lq'] = 'datasets/REDS/train_sharp_bicubic'
opt['dataroot_flow'] = None
opt['meta_info_file'] = 'basicsr/data/meta_info/meta_info_REDS_GT.txt'
opt['io_backend'] = dict(type='disk')
elif mode == 'lmdb':
opt['dataroot_gt'] = 'datasets/REDS/train_sharp_with_val.lmdb'
opt['dataroot_lq'] = 'datasets/REDS/train_sharp_bicubic_with_val.lmdb'
opt['dataroot_flow'] = None
opt['meta_info_file'] = 'basicsr/data/meta_info/meta_info_REDS_GT.txt'
opt['io_backend'] = dict(type='lmdb')
opt['val_partition'] = 'REDS4'
opt['num_frame'] = 5
opt['gt_size'] = 256
opt['interval_list'] = [1]
opt['random_reverse'] = True
opt['use_hflip'] = True
opt['use_rot'] = True
opt['use_shuffle'] = True
opt['num_worker_per_gpu'] = 1
opt['batch_size_per_gpu'] = 16
opt['scale'] = 4
opt['dataset_enlarge_ratio'] = 1
os.makedirs('tmp', exist_ok=True)
dataset = build_dataset(opt)
data_loader = build_dataloader(dataset, opt, num_gpu=0, dist=opt['dist'], sampler=None)
nrow = int(math.sqrt(opt['batch_size_per_gpu']))
padding = 2 if opt['phase'] == 'train' else 0
print('start...')
for i, data in enumerate(data_loader):
if i > 5:
break
print(i)
lq = data['lq']
gt = data['gt']
key = data['key']
print(key)
for j in range(opt['num_frame']):
torchvision.utils.save_image(
lq[:, j, :, :, :], f'tmp/lq_{i:03d}_frame{j}.png', nrow=nrow, padding=padding, normalize=False)
torchvision.utils.save_image(gt, f'tmp/gt_{i:03d}.png', nrow=nrow, padding=padding, normalize=False)
def main(mode='folder'):
"""Test paired image dataset.
Args:
mode: There are three modes: 'lmdb', 'folder', 'meta_info_file'.
"""
opt = {}
opt['dist'] = False
opt['phase'] = 'train'
opt['name'] = 'DIV2K'
opt['type'] = 'PairedImageDataset'
if mode == 'folder':
opt['dataroot_gt'] = 'datasets/DIV2K/DIV2K_train_HR_sub'
opt['dataroot_lq'] = 'datasets/DIV2K/DIV2K_train_LR_bicubic/X4_sub'
opt['filename_tmpl'] = '{}'
opt['io_backend'] = dict(type='disk')
elif mode == 'meta_info_file':
opt['dataroot_gt'] = 'datasets/DIV2K/DIV2K_train_HR_sub'
opt['dataroot_lq'] = 'datasets/DIV2K/DIV2K_train_LR_bicubic/X4_sub'
opt['meta_info_file'] = 'basicsr/data/meta_info/meta_info_DIV2K800sub_GT.txt' # noqa:E501
opt['filename_tmpl'] = '{}'
opt['io_backend'] = dict(type='disk')
elif mode == 'lmdb':
opt['dataroot_gt'] = 'datasets/DIV2K/DIV2K_train_HR_sub.lmdb'
opt['dataroot_lq'] = 'datasets/DIV2K/DIV2K_train_LR_bicubic_X4_sub.lmdb' # noqa:E501
opt['io_backend'] = dict(type='lmdb')
opt['gt_size'] = 128
opt['use_hflip'] = True
opt['use_rot'] = True
opt['use_shuffle'] = True
opt['num_worker_per_gpu'] = 2
opt['batch_size_per_gpu'] = 16
opt['scale'] = 4
opt['dataset_enlarge_ratio'] = 1
os.makedirs('tmp', exist_ok=True)
dataset = build_dataset(opt)
data_loader = build_dataloader(dataset,
opt,
num_gpu=0,
dist=opt['dist'],
sampler=None)
nrow = int(math.sqrt(opt['batch_size_per_gpu']))
padding = 2 if opt['phase'] == 'train' else 0
print('start...')
for i, data in enumerate(data_loader):
if i > 5:
break
print(i)
lq = data['lq']
gt = data['gt']
lq_path = data['lq_path']
gt_path = data['gt_path']
print(lq_path, gt_path)
torchvision.utils.save_image(lq,
f'tmp/lq_{i:03d}.png',
nrow=nrow,
padding=padding,
normalize=False)
torchvision.utils.save_image(gt,
f'tmp/gt_{i:03d}.png',
nrow=nrow,
padding=padding,
normalize=False)