详细内容可参考: Some notes about computational imaging: Calculate picture SSIM and PSNR using Python
分两种情况:
1. 在网络训练过程中计算Output和Groundtruth之间的SSIM,作为损失函数;
2. 直接计算两张图片之间的SSIM。
情况1
https://github.com/Po-Hsun-Su/pytorch-ssim
可直接使用上述代码。
情况2
参考上述代码,可将上述代码更改如下:
def ssim(img1,img2):
img1 = torch.from_numpy(np.rollaxis(img1, 2)).float().unsqueeze(0)/255.0
img2 = torch.from_numpy(np.rollaxis(img2, 2)).float().unsqueeze(0)/255.0
img1 = Variable( img1, requires_grad=False) # torch.Size([256, 256, 3])
img2 = Variable( img2, requires_grad = False)
ssim_value = pytorch_ssim.ssim(img1, img2).item()
return ssim_value
而计算psnr的代码见:Python计算图片PSNR_NLOS的博客-CSDN博客_python计算psnr
总体而言,计算两张图片的psnr和ssim的代码如下所示:
import numpy
import numpy as np
import math
import cv2
import torch
import pytorch_ssim
from torch.autograd import Variable
original = cv2.imread("1.png") # numpy.adarray
contrast = cv2.imread("2.png",1)
def psnr(img1, img2):
img1 = np.float64(img1)
img2 = np.float64(img2)
mse = numpy.mean( (img1 - img2) ** 2 )
if mse == 0:
return 100
PIXEL_MAX = 255.0
return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
def ssim(img1,img2):
img1 = torch.from_numpy(np.rollaxis(img1, 2)).float().unsqueeze(0)/255.0
img2 = torch.from_numpy(np.rollaxis(img2, 2)).float().unsqueeze(0)/255.0
img1 = Variable( img1, requires_grad=False) # torch.Size([256, 256, 3])
img2 = Variable( img2, requires_grad = False)
ssim_value = pytorch_ssim.ssim(img1, img2).item()
return ssim_value
psnrValue = psnr(original,contrast)
ssimValue = ssim(original,contrast)
print(psnrValue)
print(ssimValue)