def __init__(
self,
model='net-lin',
net='alex',
colorspace='rgb',
spatial=False,
use_gpu=True,
gpu_ids=[0],
version='0.1'
): # VGG using our perceptually-learned weights (LPIPS metric)
# def __init__(self, model='net', net='vgg', use_gpu=True): # "default" way of using VGG as a perceptual loss
super(PerceptualLoss, self).__init__()
print('Setting up Perceptual loss...')
self.use_gpu = use_gpu
self.spatial = spatial
self.gpu_ids = gpu_ids
self.model = dist_model.DistModel()
self.model.initialize(model=model,
net=net,
use_gpu=use_gpu,
colorspace=colorspace,
spatial=self.spatial,
gpu_ids=gpu_ids,
version=version)
print('...[%s] initialized' % self.model.name())
print('...Done')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
results = {'Set5': {'psnr': [], 'ssim': [], 'lpips': []},
'Set14': {'psnr': [], 'ssim': [], 'lpips': []},
'BSD100': {'psnr': [], 'ssim': [], 'lpips': []},
'Urban100': {'psnr': [], 'ssim': [], 'lpips': []},
'SunHays80': {'psnr': [], 'ssim': [], 'lpips': []}}
generator = Generator(3, filters=64, num_res_blocks=opt.residual_blocks, up_scale=opt.upSampling).eval().to(device)
checkpoint = torch.load('./checkpoints/model_G_i0085'
'00_best.pth', map_location=torch.device('cpu')) # GPU to CPU
generator.load_state_dict(checkpoint['model_G'])
test_set = TestDatasetFromFolder('data/test', upscale_factor=opt.upSampling)
test_loader = DataLoader(dataset=test_set, num_workers=1, batch_size=1, shuffle=False)
test_bar = tqdm(test_loader, desc='[testing benchmark datasets]')
model_LPIPS = dm.DistModel()
model_LPIPS.initialize(model='net-lin', net='alex', use_gpu=False)
generator.eval()
out_path = 'benchmark_results/SRF_' + str(opt.upSampling) + '/'
if not os.path.exists(out_path):
os.makedirs(out_path)
for image_name, lr_image, hr_restore_img, hr_image in test_bar:
image_name = image_name[0]
lr_image = F.interpolate(lr_image, scale_factor=opt.upSampling, mode="nearest") # Tensor [0, 1]
lr_image = lr_image.data.float().cpu().clamp_(0, 1)
sr_image, _, _ = generator(lr_image)
sr_image = sr_image.data.float().cpu().clamp_(0, 1)
# sr_image = torch.clamp(sr_image, 0, 1) # for correct display
hr_image = hr_image.data.float().cpu().clamp_(0, 1)