diff --git a/ssim/__init__.py b/ssim/__init__.py index 7d6d8e0..faf467c 100644 --- a/ssim/__init__.py +++ b/ssim/__init__.py @@ -10,7 +10,7 @@ import numpy import scipy.ndimage -from numpy.ma.core import exp +from numpy.ma.core import exp, sqrt from scipy.constants.constants import pi from compat import ImageOps @@ -30,6 +30,11 @@ def _to_grayscale(im): return gray, alpha +def convolve_gaussian_2d(image, gaussian_kernel_1d): + result = scipy.ndimage.filters.correlate1d(image, gaussian_kernel_1d, axis = 0) + result = scipy.ndimage.filters.correlate1d(result, gaussian_kernel_1d, axis = 1) + return result + def compute_ssim(im1, im2, gaussian_kernel_sigma=1.5, gaussian_kernel_width=11): """ The function to compute SSIM @@ -38,15 +43,14 @@ def compute_ssim(im1, im2, gaussian_kernel_sigma=1.5, gaussian_kernel_width=11): @return: SSIM float value """ - #Gaussian kernel definition - gaussian_kernel = numpy.zeros((gaussian_kernel_width, gaussian_kernel_width)) - + # 1D Gaussian kernel definition + gaussian_kernel_1d = numpy.ndarray((gaussian_kernel_width)) + mu = int(gaussian_kernel_width / 2) + #Fill Gaussian kernel - for i in range(gaussian_kernel_width): - for j in range(gaussian_kernel_width): - gaussian_kernel[i, j] = \ - (1 / (2 * pi * (gaussian_kernel_sigma ** 2))) * \ - exp(-(((i - 5) ** 2) + ((j - 5) ** 2)) / (2 * (gaussian_kernel_sigma ** 2))) + for i in xrange(gaussian_kernel_width): + gaussian_kernel_1d[i] = (1 / (sqrt(2 * pi) * (gaussian_kernel_sigma))) * \ + exp(-(((i - mu) ** 2)) / (2 * (gaussian_kernel_sigma ** 2))) # convert the images to grayscale img_mat_1, img_alpha_1 = _to_grayscale(im1) @@ -63,8 +67,8 @@ def compute_ssim(im1, im2, gaussian_kernel_sigma=1.5, gaussian_kernel_width=11): img_mat_12 = img_mat_1 * img_mat_2 #Means obtained by Gaussian filtering of inputs - img_mat_mu_1 = scipy.ndimage.filters.convolve(img_mat_1, gaussian_kernel) - img_mat_mu_2 = scipy.ndimage.filters.convolve(img_mat_2, gaussian_kernel) + img_mat_mu_1 = convolve_gaussian_2d(img_mat_1, gaussian_kernel_1d) + img_mat_mu_2 = convolve_gaussian_2d(img_mat_2, gaussian_kernel_1d) #Squares of means img_mat_mu_1_sq = img_mat_mu_1 ** 2 @@ -72,15 +76,15 @@ def compute_ssim(im1, im2, gaussian_kernel_sigma=1.5, gaussian_kernel_width=11): img_mat_mu_12 = img_mat_mu_1 * img_mat_mu_2 #Variances obtained by Gaussian filtering of inputs' squares - img_mat_sigma_1_sq = scipy.ndimage.filters.convolve(img_mat_1_sq, gaussian_kernel) - img_mat_sigma_2_sq = scipy.ndimage.filters.convolve(img_mat_2_sq, gaussian_kernel) + img_mat_sigma_1_sq = convolve_gaussian_2d(img_mat_1_sq, gaussian_kernel_1d) + img_mat_sigma_2_sq = convolve_gaussian_2d(img_mat_2_sq, gaussian_kernel_1d) #Covariance - img_mat_sigma_12 = scipy.ndimage.filters.convolve(img_mat_12, gaussian_kernel) + img_mat_sigma_12 = convolve_gaussian_2d(img_mat_12, gaussian_kernel_1d) #Centered squares of variances - img_mat_sigma_1_sq = img_mat_sigma_1_sq - img_mat_mu_1_sq - img_mat_sigma_2_sq = img_mat_sigma_2_sq - img_mat_mu_2_sq + img_mat_sigma_1_sq -= img_mat_mu_1_sq + img_mat_sigma_2_sq -= img_mat_mu_2_sq img_mat_sigma_12 = img_mat_sigma_12 - img_mat_mu_12 #set k1,k2 & c1,c2 to depend on L (width of color map)