def demo_test(net, test_loader, scale_factor, dataset_name, video_name):
PSNR_list = []
SSIM_list = []
with torch.no_grad():
for idx_iter, (LR, HR, SR_buicbic) in enumerate(test_loader):
LR, HR = Variable(LR).cuda(), Variable(HR).cuda()
SR = net(LR)
SR = torch.clamp(SR, 0, 1)
PSNR_list.append(psnr2(SR, HR[:, :, 3, :, :]))
SSIM_list.append(ssim(SR, HR[:, :, 3, :, :]))
if not os.path.exists('results/' + dataset + '/' + video_name):
os.makedirs('./results/' + dataset + '/' + video_name)
# ## save y images
# SR_img = transforms.ToPILImage()(SR[0, :, :, :].cpu())
# SR_img.save('results/' + dataset_name + '/' + video_name + '/sr_y_' + str(idx_iter+1).rjust(2, '0') + '.png')
#
# ## save rgb images
# SR_buicbic[:, 0, :, :] = SR[:, 0, :, :].cpu()
# SR_rgb = (ycbcr2rgb(SR_buicbic[0,:,:,:].permute(2,1,0))).permute(2,1,0)
# SR_rgb = torch.clamp(SR_rgb, 0, 1)
# SR_img = transforms.ToPILImage()(SR_rgb)
# SR_img.save('results/' + dataset_name + '/' + video_name + '/sr_rgb_' + str(idx_iter+1).rjust(2, '0') + '.png')
PSNR_mean = float(torch.cat(PSNR_list, 0)[2:-2].data.cpu().mean())
SSIM_mean = float(torch.cat(SSIM_list, 0)[2:-2].data.cpu().mean())
print(video_name + ' psnr: ' + str(PSNR_mean) + ' ssim: ' +
str(SSIM_mean))
return PSNR_mean, SSIM_mean
def cal_ssim(video1, video2, window_size=11, size_average=False):
N, C, L, H, W = np.shape(video1)
ssim_avg = 0
for i in range(N):
for j in range(L):
frame_noisy = video1[i, :, j, :, :]
frame_noisy = frame_noisy.view(1, C, H, W)
frame_gt = video2[i, :, j, :, :]
frame_gt = frame_gt.view(1, C, H, W)
ssim_frame = ssim(frame_noisy.detach(), frame_gt.detach())
ssim_avg = ssim_avg + ssim_frame.cpu().numpy()
return ssim_avg / (N * L)