File size: 2,047 Bytes
bd1c686 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | import torch
import torch.nn.functional as F
from math import log10
import cv2
import numpy as np
import torchvision
from skimage.metrics import structural_similarity as ssim
def to_psnr(frame_out, gt):
mse = F.mse_loss(frame_out, gt, reduction='none')
mse_split = torch.split(mse, 1, dim=0)
mse_list = [torch.mean(torch.squeeze(mse_split[ind])).item() for ind in range(len(mse_split))]
rmse_list = np.sqrt(mse_list) ##
intensity_max = 1.0
psnr_list = [10.0 * log10(intensity_max / mse) for mse in mse_list]
return psnr_list, rmse_list
def to_ssim_skimage(dehaze, gt):
dehaze_list = torch.split(dehaze, 1, dim=0)
gt_list = torch.split(gt, 1, dim=0)
dehaze_list_np = [dehaze_list[ind].permute(0, 2, 3, 1).data.cpu().numpy().squeeze() for ind in range(len(dehaze_list))]
gt_list_np = [gt_list[ind].permute(0, 2, 3, 1).data.cpu().numpy().squeeze() for ind in range(len(dehaze_list))]
ssim_list = [ssim(dehaze_list_np[ind], gt_list_np[ind], data_range=1, multichannel=True) for ind in range(len(dehaze_list))]
return ssim_list
def predict(gridnet, test_data_loader):
psnr_list = []
for batch_idx, (frame1, frame2, frame3) in enumerate(test_data_loader):
with torch.no_grad():
frame1 = frame1.to(torch.device('cuda'))
frame3 = frame3.to(torch.device('cuda'))
gt = frame2.to(torch.device('cuda'))
# print(frame1)
frame_out = gridnet(frame1, frame3)
# print(frame_out)
frame_debug = torch.cat((frame1, frame_out, gt, frame3), dim =0)
filepath = "./image" + str(batch_idx) + '.png'
torchvision.utils.save_image(frame_debug, filepath)
# print(frame_out)
# img = np.asarray(frame_out.cpu()).astype(float)
# cv2.imwrite(filepath , img)
# --- Calculate the average PSNR --- #
a, b = to_psnr(frame_out, gt)
psnr_list.extend(a)
avr_psnr = sum(psnr_list) / len(psnr_list)
return avr_psnr
|