def evaluate(hr: torch.tensor, sr: torch.tensor):
batch_size, _, h, w = hr.shape
psnrs, ssims, msssims = [], [], []
for b in range(batch_size):
psnrs.append(get_psnr(hr[b], sr[b]))
ssims.append(get_ssim(hr[b].unsqueeze(0), sr[b].unsqueeze(0)).item())
if h > 160 and w > 160:
msssim = get_msssim(hr[b].unsqueeze(0), sr[b].unsqueeze(0)).item()
else:
msssim = 0
msssims.append(msssim)
return np.array(psnrs).mean(), np.array(ssims).mean(), np.array(
msssims).mean()
for image in tqdm(set5):
lr, hr = get_tensor(image)
if k == 'Freq_Fusion':
lr_high, lr_low = pass_filter(lr)
sr, out_img, out_hf, out_lf = model(lr, lr_high, lr_low)
elif 'hm' in k:
lr_hf = high_pass_filter_hard_kernel(lr)
sr, _ = model(lr, lr_hf)
else:
sr, fea = model(lr)
results[image][k] = sr
evals['psnr'][image.split('/')[-1]][k] = get_psnr(sr, hr)
evals['ssim'][image.split('/')[-1]][k] = get_ssim(sr, hr).item()
evals['ms_ssim'][image.split('/')[-1]][k] = get_msssim(sr,
hr).item()
del model
torch.cuda.empty_cache()
for meas in ['psnr', 'ssim', 'ms_ssim']:
evals[meas]['mean'] = dict()
for k in weights.keys():
array = [evals[meas][image.split('/')[-1]][k] for image in set5]
evals[meas]['mean'][k] = np.mean(np.array(array))
for image in set5:
lr, hr = get_tensor(image)
rname = result_dir + image.split('/')[-1]
img = [torch.cat((lr, torch.zeros_like(lr)), dim=-2)
] + [results[image][k] for k in weights.keys()] + [hr]
imshow(img, filename=rname, visualize=False)