train_loader = DataLoader(dataset=train_set, num_workers=opt.num_workers, batch_size=opt.batch_size, shuffle=True)
val_set = loader.Val_Deresnet_Dataset(PATHS['camerasr'][opt.artifacts]['valid_hr'],
lr_dir=PATHS['camerasr'][opt.artifacts]['valid_lr'], **vars(opt))
val_loader = DataLoader(dataset=val_set, num_workers=1, batch_size=1, shuffle=False)
else:
train_set = loader.TrainDataset(PATHS[opt.dataset][opt.artifacts]['hr']['train'], cropped=True, **vars(opt))
train_loader = DataLoader(dataset=train_set, num_workers=opt.num_workers, batch_size=opt.batch_size, shuffle=True)
val_set = loader.ValDataset(PATHS[opt.dataset][opt.artifacts]['hr']['valid'],
lr_dir=PATHS[opt.dataset][opt.artifacts]['lr']['valid'], **vars(opt))
val_loader = DataLoader(dataset=val_set, num_workers=1, batch_size=1, shuffle=False)
# prepare neural networks
if opt.generator.lower() == 'dsgan':
model_g = model.Generator(n_res_blocks=opt.num_res_blocks)
elif opt.generator.lower() == 'deresnet':
model_g = model.De_resnet(n_res_blocks=opt.num_res_blocks, scale=opt.upscale_factor)
else:
raise NotImplementedError('Generator model [{:s}] not recognized'.format(opt.generator))
print('# Initializing {}'.format(opt.generator))
print('# generator parameters:', sum(param.numel() for param in model_g.parameters()), '\n')
model_d = model.Discriminator(kernel_size=opt.kernel_size, wgan=opt.wgan, highpass=opt.highpass,
D_arch=opt.discriminator, norm_layer=opt.norm_layer, filter_type=opt.filter,
cs=opt.cat_or_sum)
print('# discriminator parameters:', sum(param.numel() for param in model_d.parameters()))
g_loss_module = loss.GeneratorLoss(**vars(opt))
# filters are used for generating validation images
filter_low_module = model.FilterLow(kernel_size=opt.kernel_size, gaussian=opt.filter == 'gau', include_pad=False)
if not os.path.exists(sdsr_lr_dir):
os.makedirs(sdsr_lr_dir)
if not os.path.exists(tdsr_hr_dir):
os.makedirs(tdsr_hr_dir)
if not os.path.exists(tdsr_lr_dir):
os.makedirs(tdsr_lr_dir)
if not os.path.exists(tdsr_lr_img_dir):
os.makedirs(tdsr_lr_img_dir)
if not os.path.exists(tdsr_lr_ddm_dir):
os.makedirs(tdsr_lr_ddm_dir)
# prepare neural networks
if opt.Generator == 'DSGAN':
model_g = model.Generator(n_res_blocks=opt.num_res_blocks)
elif opt.Generator == 'DeResnet':
model_g = model.De_resnet(n_res_blocks=opt.num_res_blocks)
if opt.Discriminator == 'FSD':
# model_d = model.Discriminator_wavelet(patchgan=False)
model_d = model.Discriminator_Gau(patchgan=False, gaussian=True)
convnet = [[5, 1, 2], [5, 1, 2], [5, 1, 2], [5, 1, 2]]
elif opt.Discriminator == 'n_layer_D_s1':
model_d = model.Discriminator_wavelet(cs='cat', patchgan='s1')
convnet = [[4, 1, 1], [4, 1, 1], [4, 1, 1], [4, 1, 1]]
elif opt.Discriminator == 'n_layer_D_s2':
model_d = model.Discriminator_wavelet(cs='cat', patchgan='s2')
convnet = [[4, 2, 1], [4, 2, 1], [4, 1, 1], [4, 1, 1]]
model_g = model_g.eval()
model_d = model_d.eval()
print('# generator parameters:',
sum(param.numel() for param in model_g.parameters()))