Merge pull request #7 from koteth/master

refactoring + autoresize + multiple comparison

setup.py

@@ -32,7 +32,7 @@
     url = "https://github.com/jterrace/pyssim",
     entry_points = {
         'console_scripts':[
-            'pyssim = ssim.__main__:main'
+            'pyssim = ssim.ssim:main'
         ]
     },
     packages = find_packages()

ssim/__init__.py

@@ -1,108 +1,12 @@
-"""
-This module computes the Structured Similarity Image Metric (SSIM)
+from ssim import get_gaussian_kernel, SSIM, SSIMImage
+from compat import Image
 
-Created on 21 nov. 2011
-@author: Antoine Vacavant, ISIT lab, [email protected], http://isit.u-clermont1.fr/~anvacava
-
-Modified by Christopher Godfrey, on 17 July 2012 (lines 32-34)
-Modified by Jeff Terrace, starting 29 August 2012
-"""
-
-import numpy
-import scipy.ndimage
-from numpy.ma.core import exp, sqrt
-from scipy.constants.constants import pi
-from compat import ImageOps
-
-def _to_grayscale(im):
-    """
-    Convert PIL image to numpy grayscale array and numpy alpha array
-    @param im: PIL Image object
-    @return (gray, alpha), both numpy arrays
-    """
-    gray = numpy.asarray(ImageOps.grayscale(im)).astype(numpy.float)
-
-    imbands = im.getbands()
-    if 'A' in imbands:
-        alpha = numpy.asarray(im.split()[-1]).astype(numpy.float)
-    else:
-        alpha = None
-
-    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):
+def compute_ssim(image1, image2, gaussian_kernel_sigma=1.5, gaussian_kernel_width=11):
     """
     The function to compute SSIM
     @param im1: PIL Image object
     @param im2: PIL Image object
     @return: SSIM float value
     """
-
-    # 1D Gaussian kernel definition
-    gaussian_kernel_1d = numpy.ndarray((gaussian_kernel_width))
-    mu = int(gaussian_kernel_width / 2)
-
-    #Fill Gaussian kernel
-    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)
-    img_mat_2, img_alpha_2 = _to_grayscale(im2)
-
-    # don't count pixels where both images are both fully transparent
-    if img_alpha_1 is not None and img_alpha_2 is not None:
-        img_mat_1[img_alpha_1 == 255] = 0
-        img_mat_2[img_alpha_2 == 255] = 0
-
-    #Squares of input matrices
-    img_mat_1_sq = img_mat_1 ** 2
-    img_mat_2_sq = img_mat_2 ** 2
-    img_mat_12 = img_mat_1 * img_mat_2
-
-    #Means obtained by Gaussian filtering of inputs
-    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
-    img_mat_mu_2_sq = img_mat_mu_2 ** 2
-    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 = 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 = convolve_gaussian_2d(img_mat_12, gaussian_kernel_1d)
-
-    #Centered squares of variances
-    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)
-    l = 255
-    k_1 = 0.01
-    c_1 = (k_1 * l) ** 2
-    k_2 = 0.03
-    c_2 = (k_2 * l) ** 2
-
-    #Numerator of SSIM
-    num_ssim = (2 * img_mat_mu_12 + c_1) * (2 * img_mat_sigma_12 + c_2)
-
-    #Denominator of SSIM
-    den_ssim = (img_mat_mu_1_sq + img_mat_mu_2_sq + c_1) * \
-               (img_mat_sigma_1_sq + img_mat_sigma_2_sq + c_2)
-
-    #SSIM
-    ssim_map = num_ssim / den_ssim
-    index = numpy.average(ssim_map)
-
-    return index
+    gaussian_kernel_1d = get_gaussian_kernel(gaussian_kernel_width, gaussian_kernel_sigma)
+    return SSIM(image1, gaussian_kernel_1d).ssim_value(image2)

ssim/__main__.py

@@ -1,18 +1,4 @@
-import argparse
-import ssim
-from ssim.compat import Image
+from ssim import main
 
-def main():
-    parser = argparse.ArgumentParser(prog="pyssim",
-                                     description="Compares two images using the SSIM metric")
-    parser.add_argument('base_image', metavar='image1.png', type=argparse.FileType('r'))
-    parser.add_argument('comparison_image', metavar='image2.png', type=argparse.FileType('r'))
-    args = parser.parse_args()
-
-    im1 = Image.open(args.base_image)
-    im2 = Image.open(args.comparison_image)
-
-    print ssim.compute_ssim(im1, im2)
-
 if __name__ == '__main__':
     main()

ssim/ssim.py

@@ -0,0 +1,111 @@
+import argparse
+import glob
+from compat import Image
+import numpy
+import scipy.ndimage
+from numpy.ma.core import exp, sqrt
+from scipy.constants.constants import pi
+from utils import to_grayscale, convolve_gaussian_2d, get_gaussian_kernel
+from argparse import RawTextHelpFormatter
+
+class SSIMImage(object):
+    """SSIMImage 
+    Parameters:
+      img: Image or file name
+      gaussian_kernel_1d
+      size: new images size
+
+    Attributes:
+      img: PIL original Image.
+      img_gray: grayscale Image.
+      img_gray_squared: squared img_gray
+      img_gray_mu: img_gray convolved with gaussian kernel
+      img_gray_mu_squared: squared img_gray_mu
+      img_gray_sigma_squared: img_gray convolved with gaussian kernel - img_gray_mu_squared
+
+    """
+    def __init__(self, img, gaussian_kernel_1d, size = None):
+        self.gaussian_kernel_1d = gaussian_kernel_1d
+        if isinstance(img, basestring):
+            self.img = Image.open(img)
+        else:
+            self.img = img
+        if size:
+            self.img = self.img.resize(size)
+        self.size = self.img.size
+        self.img_gray, self.img_alpha = to_grayscale(self.img)
+        if self.img_alpha is not None:
+            self.img_gray[self.img_alpha == 255] = 0
+        self.img_gray_squared = self.img_gray ** 2
+        self.img_gray_mu = convolve_gaussian_2d(self.img_gray,  self.gaussian_kernel_1d)
+        self.img_gray_mu_squared = self.img_gray_mu ** 2
+        self.img_gray_sigma_squared = convolve_gaussian_2d(self.img_gray_squared, self.gaussian_kernel_1d)
+        self.img_gray_sigma_squared -= self.img_gray_mu_squared
+
+class SSIM(object):
+    """
+    SSIM: the class comupute SSIM between two images 
+    Parameters:
+      img1: reference image
+      gaussian_kernel_1d
+      l, k_1, k_2
+    """
+    def __init__(self, img1, gaussian_kernel_1d, l =255, k_1 =0.01, k_2 = 0.03):
+        #set k1,k2 & c1,c2 to depend on L (width of color map)
+        self.c_1 = (k_1 * l) ** 2
+        self.c_2 = (k_2 * l) ** 2
+        self.gaussian_kernel_1d = gaussian_kernel_1d
+        self.img1 = SSIMImage(img1, gaussian_kernel_1d)
+
+    def ssim_value(self, img2):
+        """
+        The function return SSIM value from the reference image and the 
+            input image
+        Parameters:
+          img2: input image 
+          return: SSIM float value
+        """
+        self.img2 = SSIMImage(img2, self.gaussian_kernel_1d,self.img1.size)
+        self.img_mat_12 = self.img1.img_gray * self.img2.img_gray
+        self.img_mat_sigma_12 = convolve_gaussian_2d(self.img_mat_12, self.gaussian_kernel_1d)
+        self.img_mat_mu_12 = self.img1.img_gray_mu * self.img2.img_gray_mu
+        self.img_mat_sigma_12 = self.img_mat_sigma_12 - self.img_mat_mu_12
+        #Numerator of SSIM
+        num_ssim = (2 * self.img_mat_mu_12 + self.c_1) * (2 * self.img_mat_sigma_12 + self.c_2)
+        #Denominator of SSIM
+        den_ssim = (self.img1.img_gray_mu_squared + self.img2.img_gray_mu_squared + self.c_1) * \
+                   (self.img1.img_gray_sigma_squared + self.img2.img_gray_sigma_squared + self.c_2)
+        #SSIM
+        ssim_map = num_ssim / den_ssim
+        index = numpy.average(ssim_map)
+        return index
+
+def main():
+    description = "\n".join([ "Compares an image with a list of images using the SSIM metric",
+                            " example:", '  pyssim test-images/test1-1.png "test-images/*"'])
+
+    parser = argparse.ArgumentParser(prog="pyssim", formatter_class=RawTextHelpFormatter,
+                                     description=description)
+    parser.add_argument('base_image', metavar='image1.png', type=argparse.FileType('r'))
+    parser.add_argument('comparison_images', metavar='image path with* or image2.png')
+    args = parser.parse_args()
+    gaussian_kernel_sigma=1.5
+    gaussian_kernel_width=11
+    gaussian_kernel_1d = get_gaussian_kernel(gaussian_kernel_width, gaussian_kernel_sigma)
+
+    comparison_images = glob.glob(args.comparison_images)
+    is_a_single_image = len(comparison_images) == 1
+
+    for comparison_image in comparison_images:
+        try:
+            ssim_value = SSIM(args.base_image.name, gaussian_kernel_1d).ssim_value(comparison_image)
+            if is_a_single_image:
+                print ssim_value
+            else:
+                print "%s - %s: %s" % (args.base_image.name, comparison_image, ssim_value)
+
+        except Exception, e:
+            print e
+
+if __name__ == '__main__':
+    main()

ssim/utils.py

@@ -0,0 +1,38 @@
+from compat import Image
+import numpy
+import scipy.ndimage
+from numpy.ma.core import exp, sqrt
+from scipy.constants.constants import pi
+from compat import ImageOps
+
+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 get_gaussian_kernel(gaussian_kernel_width=11, gaussian_kernel_sigma=1.5 ):
+    # 1D Gaussian kernel definition
+    gaussian_kernel_1d = numpy.ndarray((gaussian_kernel_width))
+    mu = int(gaussian_kernel_width / 2)
+
+    #Fill Gaussian kernel
+    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)))
+    return gaussian_kernel_1d
+
+def to_grayscale(im):
+    """
+    Convert PIL image to numpy grayscale array and numpy alpha array
+    @param im: PIL Image object
+    @return (gray, alpha), both numpy arrays
+    """
+    gray = numpy.asarray(ImageOps.grayscale(im)).astype(numpy.float)
+
+    imbands = im.getbands()
+    if 'A' in imbands:
+        alpha = numpy.asarray(im.split()[-1]).astype(numpy.float)
+    else:
+        alpha = None
+
+    return gray, alpha