110 lines
3.2 KiB
Python
Executable File
110 lines
3.2 KiB
Python
Executable File
import numpy as np
|
||
import imageio
|
||
from skimage.transform import resize
|
||
from scipy.ndimage import gaussian_filter
|
||
|
||
|
||
def mssim(
|
||
x: np.ndarray,
|
||
y: np.ndarray,
|
||
) -> float:
|
||
# Standard choice for the parameters
|
||
K1 = 0.01
|
||
K2 = 0.03
|
||
sigma = 1.5
|
||
truncate = 3.5
|
||
m = 1
|
||
C1 = (K1 * m) ** 2
|
||
C2 = (K2 * m) ** 2
|
||
|
||
# radius size of the local window (needed for
|
||
# normalizing the standard deviation)
|
||
r = int(truncate * sigma + 0.5)
|
||
win_size = 2 * r + 1
|
||
# use these arguments for the gaussian filtering
|
||
# e.g. filtered = gaussian_filter(x, **filter_args)
|
||
filter_args = {
|
||
'sigma': sigma,
|
||
'truncate': truncate
|
||
}
|
||
|
||
# Implement Eq. (9) from assignment sheet
|
||
# S should be an "image" of the SSIM evaluated for a window
|
||
# centered around the corresponding pixel in the original input image
|
||
S = np.ones_like(x)
|
||
|
||
mu_x = gaussian_filter(x, **filter_args)
|
||
mu_y = gaussian_filter(y, **filter_args)
|
||
|
||
# Compute local variances and covariance
|
||
n = win_size ** 2 # number of pixels in the window
|
||
# n~ = n/(n-1)
|
||
# E[x^2] −E[x]^2). = (gaussian_filter(x * x, **filter_args) - mu_x * mu_x)
|
||
sigma_x_sq = (gaussian_filter(x * x, **filter_args) - mu_x * mu_x) * (n / (n - 1))
|
||
sigma_y_sq = (gaussian_filter(y * y, **filter_args) - mu_y * mu_y) * (n / (n - 1))
|
||
sigma_xy = (gaussian_filter(x * y, **filter_args) - mu_x * mu_y) * (n / (n - 1))
|
||
|
||
# Compute SSIM
|
||
S = ((2 * mu_x * mu_y + C1) * (2 * sigma_xy + C2)) / \
|
||
((mu_x ** 2 + mu_y ** 2 + C1) * (sigma_x_sq + sigma_y_sq + C2))
|
||
|
||
# crop to remove boundary artifacts, return MSSIM
|
||
pad = (win_size - 1) // 2
|
||
return S[pad:-pad, pad:-pad].mean()
|
||
|
||
|
||
def psnr(
|
||
x: np.ndarray,
|
||
y: np.ndarray,
|
||
) -> float:
|
||
# Implement Eq. (2) without for loops
|
||
mse = np.mean((x - y) ** 2)
|
||
m = 1
|
||
psnr_value = 10 * np.log10(m ** 2 / mse)
|
||
return psnr_value
|
||
|
||
|
||
def psnr_for(
|
||
x: np.ndarray,
|
||
y: np.ndarray,
|
||
) -> float:
|
||
# Implement Eq. (2) using for loops
|
||
m, n = x.shape
|
||
mse = 0
|
||
# Eq. (1) from the assignment sheet for mse
|
||
for i in range(m):
|
||
for j in range(n):
|
||
mse += (x[i, j] - y[i, j]) ** 2
|
||
mse /= (m * n)
|
||
m = 1
|
||
psnr_value = 10 * np.log10(m ** 2 / mse)
|
||
return psnr_value
|
||
|
||
|
||
def interpolation_error():
|
||
x = imageio.imread('./girl.png') / 255.
|
||
shape_lower = (x.shape[0] // 2, x.shape[1] // 2)
|
||
# downsample image to half the resolution
|
||
# and successively upsample to the original resolution
|
||
# using no nearest neighbor, linear and cubic interpolation
|
||
nearest, linear, cubic = [
|
||
resize(resize(
|
||
x, shape_lower, order=order, anti_aliasing=False
|
||
), x.shape, order=order, anti_aliasing=False)
|
||
for order in [0, 1, 3]
|
||
]
|
||
|
||
for label, rescaled in zip(
|
||
['nearest', 'linear', 'cubic'],
|
||
[nearest, linear, cubic]
|
||
):
|
||
print(label)
|
||
print(mssim(x, rescaled))
|
||
print(psnr(x, rescaled))
|
||
print(psnr_for(x, rescaled))
|
||
imageio.imwrite('girl_' + label + '.png', (rescaled * 255).astype(np.uint8))
|
||
|
||
|
||
if __name__ == '__main__':
|
||
interpolation_error()
|