python设计PSNR和SSIM计算函数

PSNR

import numpy as np
import math

def PSNR(x, y):
    MSE = np.mean((x/255. - y/255.)**2)
    if MSE<1.0e-10:
        return 100
    MAX = 1
    return 20*math.log10(MAX/math.sqrt(MSE))

SSIM

import numpy as np
from PIL import Image
from scipy.signal import convolve2d

def gauss2D(shape=(3, 3), sigma=0.5):

    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
    h_sum = h.sum()
    if h_sum != 0:
        h /= h_sum
    return h

def filter2(x, kernel, mode='same'):
    return convolve2d(x, np.rot90(kernel, 2), mode=mode)

def compute_ssim(image1, image2, k1=0.01, k2=0.03, win_size=11, L=255):

    if not image1.shape == image2.shape:
        raise ValueError("输入的两张图片大小应该一样")
    if len(image1.shape) > 2:
        raise ValueError("输入的图片应为灰度图")

    M, N = image1.shape
    C1 = (k1 * L) ** 2
    C2 = (k2 * L) ** 2
    window = gauss2D(shape=(win_size, win_size), sigma=1.5)
    window = window / np.sum(np.sum(window))

    if image1.dtype == np.uint8:
        image1 = np.double(image1)
    if image2.dtype == np.uint8:
        image2 = np.double(image2)

    mu1 = filter2(image1, window, 'valid')
    mu2 = filter2(image2, window, 'valid')
    mu1_sq = mu1 * mu1
    mu2_sq = mu2 * mu2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = filter2(image1 * image1, window, 'valid') - mu1_sq
    sigma2_sq = filter2(image2 * image2, window, 'valid') - mu2_sq
    sigmal2 = filter2(image1 * image2, 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))

其他

仅做学习记录用，具体来源忘记了，侵删。