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))