## 2020年3月13日星期五

### Calculate picture SSIM and PSNR using Python

This article mainly shows the calculation of two pictures psnr and ssim using Python.

## 1. Calculation of SSIM

There are two cases:

A. Calculate the SSIM between Output and Groundtruth as a loss function during network training;

B. Calculate SSIM directly between two pictures.

### CaseA:

https://github.com/congyucn/pytorch-ssim
There is code that can be used directly on github, such as the above code.

### CaseB:

Referring to the above code, the code can be changed as follows:

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

## 2. Caculation of PSNR

https://blog.csdn.net/qazwsxrx/article/details/104550550

The code for calculating psnr can be found through the above URL.

## 3. Conclusion

Overall, the code for calculating the psnr and ssim of two pictures is as follows:

import numpy
import numpy as np
import math
import cv2
import torch
import pytorch_ssim

def psnr(img1, 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)