python PSNR代码

图像恢复常用psnr做为评价指标。psnr值越高则表示两幅图像越相似。

由于图像处理时常有两种类型的格式，一种是tensor类型，一种是numpy类型。所以这里准备了两个代码。

import torch
import numpy as np
import cv2
from torchvision import transforms

def torchPSNR(tar_img, prd_img):
    imdff = torch.clamp(prd_img,0,1) - torch.clamp(tar_img,0,1)
    rmse = (imdff**2).mean().sqrt()
    ps = 20*torch.log10(1/rmse)
    return ps

def numpyPSNR(tar_img, prd_img):
    ##这里通常会去除边框
    imdff = np.float32(prd_img) - np.float32(tar_img)
    rmse = np.sqrt(np.mean(imdff**2))
    ps = 20*np.log10(255/rmse)
    return ps

img1=cv2.imread('./2.png')
# img1=cv2.cvtColor(img1, cv2.COLOR_BGR2YCR_CB)###改变颜色
# b1,g1,r1=cv2.split(img1)  ##通道分离
img2=cv2.imread('./3.png')
# img2=cv2.cvtColor(img2, cv2.COLOR_BGR2YCR_CB)
# b2,g2,r2=cv2.split(img2)
print(numpyPSNR(img1,img2))

img1=transforms.ToTensor()(img1)
img2=transforms.ToTensor()(img2)

print(torchPSNR(img1,img2))
# print(numpyPSNR(g1,g2))
# print(numpyPSNR(b1,b2))
# print(numpyPSNR(r1,r2))

补充：

当图片的数据类型是np.uint8时。

from skimage.measure import compare_ssim,compare_psnr
psnr = compare_psnr(img1, img2, 255)