批量计算PSNR,SSIM,MSE

批量计算PSNR,SSIM,MSE

这里有几点要注意的：

1.对比两张图片的函数SSIM具体是structural_similarity：

      Ⅰ：第一个参数和第二个参数是要对比的两张图片；

      Ⅱ：参数中有个full，默认为True，如果为True会把两张图的相似部分返回，如果为False就不返回了；

      Ⅳ：返回值full为False的时候有2个，一个是mssim浮点型的相似度，1代表2张图一模一样。0代表完全不一样，grad返回两张图片不同的地方类型为ndarray；

2.SSIM返回的ndarray里面的值为[0, 1]的float型，而OpenCV的[0, 255]为uint8型，可以通过这样的方式转换：

      Ⅰ：diff = (diff * 255).astype("uint8")

3.cv2.threshold函数，设置一个阈值和一个颜色，将传入的图片没有达到这个阈值的，转换为这个颜色。

import os
import cv2
import math
import time
import numpy as np
from PIL import Image
import torch.nn.functional as F
from skimage.metrics import structural_similarity

start = time.time()

def psnr(img1, img2):  # 但是其值不能很好地反映人眼主观感受。一般取值范围：20-40.值越大，视频质量越好
    mse = np.mean((img1 / 1. - img2 / 1.) ** 2)
    if mse < 1.0e-10:
        return 100 * 1.0
    return 10 * math.log10(255.0 * 255.0 / mse)


def mse(img1, img2):  # 指参数估计值与参数真值之差平方的期望值，在图像处理算法中就是处理后图像像素值与原始像素值之差平方和的均值，越小越好
    mse = np.mean((img1 / 1. - img2 / 1.) ** 2)
    return mse


def ssim(y_true, y_pred):  # 其值可以较好地反映人眼主观感受。一般取值范围：0-1.值越大，视频质量越好
    ssim = structural_similarity(y_true, y_pred, multichannel=True, full=False)  # 如果输入是多通道（彩色）图像，设置 multichannel=True
    return ssim


path1 = './images/test1/'  # 指定输出结果文件夹
path2 = './images/test2/'  # 指定原图文件夹
f_nums = len(os.listdir(path1)) - 2
list_psnr = []
list_ssim = []
list_mse = []
file = open(r'psnr-ssim-scene3-70-100.txt', mode='w',encoding='utf-8')
for i in range(1, f_nums):
    print('第%s张图片' % i)
    img_a = Image.open(path1 + str(i) + '.png')
    img_b = Image.open(path2 + str(i) + '.png')
    img_a = np.array(img_a)
    img_b = np.array(img_b)

    psnr_num = psnr(img_a, img_b)
    ssim_num = ssim(img_a, img_b)
    mse_num = mse(img_a, img_b)
    print('psnr_num', psnr_num)
    print('ssim_num', ssim_num)
    print('mse_num', mse_num)
    list_ssim.append(ssim_num)
    list_psnr.append(psnr_num)
    list_mse.append(mse_num)
    file.write('第%s张图片:' % i + '\n')
    file.write('psnr_num, {:.5f}'.format(psnr_num) + '\n')
    file.write('ssim_num, {:.5f}'.format(ssim_num) + '\n')
    file.write('mse_num, {:.5f}'.format(mse_num) + '\n')


print("平均PSNR:", np.mean(list_psnr))  # ,list_psnr)
print("平均SSIM:", np.mean(list_ssim))  # ,list_ssim)
print("平均MSE:", np.mean(list_mse))  # ,list_mse)

elapsed = (time.time() - start)
print("Time used:", elapsed)
file.write("\n")
file.write("汇总:\n")
file.write('平均PSNR, {:.5f}'.format(np.mean(list_psnr)) + '\n')
file.write('平均SSIM, {:.5f}'.format(np.mean(list_ssim)) + '\n')
file.write('平均MSE, {:.5f}'.format(np.mean(list_mse)) + '\n')
file.write('Time used, {:.5f}'.format(elapsed) + '\n')
file.close()