import numpy as np
import math
from scipy.signal import convolve2d
import cv2
__all__ = [
"set_requires_grad", "compute_psnr", 'compute_ssim'
]
# 反向传播开关
def set_requires_grad(parameters, requires_grad=False):
# 注意这里默认是False
if not isinstance(parameters, list):
parameters = [parameters] # 包装一下(对传入的单个参数)
for param in parameters:
if param is not None:
param.requires_grad = requires_grad
def compute_psnr(img_src, img_guess):
img1_array = np.array(img_src)
img2_array = np.array(img_guess)
print(img1_array.dtype)
print(img2_array.dtype)
mse = np.mean((img1_array - img2_array)**2)
PIXEL_MAX = 255.0
if mse == 0:
return float('inf')
else:
return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
def matlab_style_gauss2D(shape=(3,3),sigma=0.5):
"""
2D gaussian mask - should give the same result as MATLAB's
fspecial('gaussian',[shape],[sigma])
"""
m,n = [(ss-1.)/2. for ss in shape]
y,x = np.ogrid[-m:m+1,-n:n+1]
h = np.exp( -(x*x + y*y) / (2.*sigma*sigma) )
h[ h < np.finfo(h.dtype).eps*h.max() ] = 0
sumh = h.sum()
if sumh != 0:
h /= sumh
return h
def filter2(x, kernel, mode='same'):
return convolve2d(x, np.rot90(kernel, 2), mode=mode)
def compute_ssim(img1, img2, k1=0.01, k2=0.03, win_size=11, L=255):
# img转array
img1_array = np.array(img1, dtype=np.uint8)
img2_array = np.array(img2, dtype=np.uint8)
print(img1_array.shape)
# 转灰度图
gray_img1 = cv2.cvtColor(img1_array, cv2.COLOR_BGR2GRAY)
gray_img2 = cv2.cvtColor(img2_array, cv2.COLOR_BGR2GRAY)
print(gray_img1.shape)
if not gray_img1.shape == gray_img2.shape:
raise ValueError("Input Imagees must have the same dimensions")
if len(gray_img1.shape) > 2:
raise ValueError("Please input the images with 1 channel")
M, N = gray_img1.shape
C1 = (k1 * L) ** 2
C2 = (k2 * L) ** 2
window = matlab_style_gauss2D(shape=(win_size, win_size), sigma=1.5)
window = window / np.sum(np.sum(window))
if gray_img1.dtype == np.uint8:
gray_img1 = np.double(gray_img1)
if gray_img2.dtype == np.uint8:
gray_img2 = np.double(gray_img2)
mu1 = filter2(gray_img1, window, 'valid')
mu2 = filter2(gray_img2, window, 'valid')
mu1_sq = mu1 * mu1
mu2_sq = mu2 * mu2
mu1_mu2 = mu1 * mu2
sigma1_sq = filter2(gray_img1 * gray_img1, window, 'valid') - mu1_sq
sigma2_sq = filter2(gray_img2 * gray_img2, window, 'valid') - mu2_sq
sigmal2 = filter2(gray_img1 * gray_img2, window, 'valid') - mu1_mu2
ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigmal2 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
return np.mean(np.mean(ssim_map))
python计算图像间PSNR和SSIM代码
于 2023-08-21 15:28:23 首次发布
![](https://img-home.csdnimg.cn/images/20240711042549.png)