-
Notifications
You must be signed in to change notification settings - Fork 8
/
utils_test.py
28 lines (23 loc) · 1.13 KB
/
utils_test.py
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
import torch
import torch.nn.functional as F
from math import log10
from skimage.metrics import structural_similarity as ssim
def to_psnr(dehaze, gt):
mse = F.mse_loss(dehaze, 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))]
intensity_max = 1.0
psnr_list = [10.0 * log10(intensity_max / mse) for mse in mse_list]
return psnr_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 to_rmse(dehaze, gt):
mse = torch.mean((dehaze - gt) ** 2, dim=[1, 2, 3])
rmse = torch.sqrt(mse)
rmse_list = [value.item() for value in rmse]
return rmse_list