File size: 3,174 Bytes
8ad2fa0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import numpy as np
from scipy import signal
from scipy import ndimage
def fspecial_gauss(size, sigma):
x, y = np.mgrid[-size // 2 + 1:size // 2 + 1, -size // 2 + 1:size // 2 + 1]
g = np.exp(-((x**2 + y**2) / (2.0 * sigma**2)))
return g / g.sum()
def calc_ssim(img1, img2, data_range=255):
img1 = img1.astype(np.float64)
img2 = img2.astype(np.float64)
size = 11
sigma = 1.5
window = fspecial_gauss(size, sigma)
K1 = 0.01
K2 = 0.03
C1 = (K1 * data_range)**2
C2 = (K2 * data_range)**2
mu1 = signal.fftconvolve(window, img1, mode='valid')
mu2 = signal.fftconvolve(window, img2, mode='valid')
mu1_sq = mu1 * mu1
mu2_sq = mu2 * mu2
mu1_mu2 = mu1 * mu2
sigma1_sq = signal.fftconvolve(window, img1 * img1, mode='valid') - mu1_sq
sigma2_sq = signal.fftconvolve(window, img2 * img2, mode='valid') - mu2_sq
sigma12 = signal.fftconvolve(window, img1 * img2, mode='valid') - mu1_mu2
return (((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
(sigma1_sq + sigma2_sq + C2)),
(2.0 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2))
def calc_msssim(img1, img2, data_range=255):
'''
img1 and img2 are 2D arrays
'''
level = 5
weight = np.array([0.0448, 0.2856, 0.3001, 0.2363, 0.1333])
height, width = img1.shape
if height < 176 or width < 176:
# according to HM implementation
level = 4
weight = np.array([0.0517, 0.3295, 0.3462, 0.2726])
if height < 88 or width < 88:
assert False
downsample_filter = np.ones((2, 2)) / 4.0
im1 = img1.astype(np.float64)
im2 = img2.astype(np.float64)
mssim = np.array([])
mcs = np.array([])
for _ in range(level):
ssim_map, cs_map = calc_ssim(im1, im2, data_range=data_range)
mssim = np.append(mssim, ssim_map.mean())
mcs = np.append(mcs, cs_map.mean())
filtered_im1 = ndimage.filters.convolve(im1, downsample_filter,
mode='reflect')
filtered_im2 = ndimage.filters.convolve(im2, downsample_filter,
mode='reflect')
im1 = filtered_im1[::2, ::2]
im2 = filtered_im2[::2, ::2]
return (np.prod(mcs[0:level - 1]**weight[0:level - 1]) *
(mssim[level - 1]**weight[level - 1]))
def calc_msssim_rgb(img1, img2, data_range=255):
'''
img1 and img2 are arrays with 3xHxW
'''
msssim = 0
for i in range(3):
msssim += calc_msssim(img1[i, :, :], img2[i, :, :], data_range)
return msssim / 3
def calc_psnr(img1, img2, data_range=255):
'''
img1 and img2 are arrays with same shape
'''
img1 = img1.astype(np.float64)
img2 = img2.astype(np.float64)
mse = np.mean(np.square(img1 - img2))
if np.isnan(mse) or np.isinf(mse):
return -999.9
if mse > 1e-10:
psnr = 10 * np.log10(data_range * data_range / mse)
else:
psnr = 999.9
if psnr > 99.9:
psnr = 99.9
return psnr
|