| 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')) |
| |
|
|
| frame_out = gridnet(frame1, frame3) |
| |
| frame_debug = torch.cat((frame1, frame_out, gt, frame3), dim =0) |
| filepath = "./image" + str(batch_idx) + '.png' |
| torchvision.utils.save_image(frame_debug, filepath) |
| |
| |
| |
| |
|
|
|
|
|
|
| |
| a, b = to_psnr(frame_out, gt) |
| psnr_list.extend(a) |
| avr_psnr = sum(psnr_list) / len(psnr_list) |
| return avr_psnr |
|
|
|
|
|
|
