structural_similarity: add option to return full ssim image in addition to the mean

skimage/measure/_structural_similarity.py

@@ -76,7 +76,7 @@ def _discard_edges(X, pad):
 
 def structural_similarity(X, Y, win_size=None, gradient=False,
                           dynamic_range=None, multichannel=None,
-                          gaussian_weights=False):
+                          gaussian_weights=False, full=False):
     """Compute the mean structural similarity index between two images.
 
     Parameters
@@ -99,6 +99,9 @@ def structural_similarity(X, Y, win_size=None, gradient=False,
     gaussian_weights : bool
         If True, each patch (of size win_size) has its mean and variance
         spatially weighted by a normalized Gaussian kernel of width sigma=1.5.
+    full : bool
+        If True, return the full structural similarity image instead of the
+        mean value
 
     Returns
     -------
@@ -107,6 +110,8 @@ def structural_similarity(X, Y, win_size=None, gradient=False,
     grad : ndarray
         Gradient of the structural similarity index between X and Y [2]. This
         is only returned if `gradient` is set to True.
+    S : ndarray
+        Full SSIM image.  This is only returned if `full` is set to True.
 
     Notes
     -----
@@ -147,16 +152,25 @@ def structural_similarity(X, Y, win_size=None, gradient=False,
         mssim = np.empty(nch)
         if gradient:
             G = np.empty(X.shape)
+        if full:
+            S = np.empty(X.shape)
         for ch in range(nch):
-            if gradient:
-                mssim[..., ch], G[..., ch] = structural_similarity(
-                    X[..., ch], Y[..., ch], **args)
+            ch_result = structural_similarity(X[..., ch], Y[..., ch], **args)
+            if gradient and full:
+                mssim[..., ch], G[..., ch], S[..., ch] = ch_result
+            elif gradient:
+                mssim[..., ch], G[..., ch] = ch_result
+            elif full:
+                mssim[..., ch], S[..., ch] = ch_result
             else:
-                mssim[..., ch] = structural_similarity(
-                    X[..., ch], Y[..., ch], **args)
+                mssim[..., ch] = ch_result
         mssim = mssim.mean()
-        if gradient:
+        if gradient and full:
+            return mssim, G, S
+        elif gradient:
             return mssim, G
+        elif full:
+            return mssim, S
         else:
             return mssim
 
@@ -230,6 +244,14 @@ def structural_similarity(X, Y, win_size=None, gradient=False,
         grad += filter_func((ux * (A2 - A1) - uy * (B2 - B1) * S) / D,
                             **filter_args)
         grad *= (2 / X.size)
-        return mssim, grad
+
+        if full:
+            return mssim, grad, S
+        else:
+            return mssim, grad
     else:
-        return mssim
+        if full:
+            return mssim, S
+        else:
+            return mssim
+

skimage/measure/tests/test_structural_similarity.py

@@ -41,6 +41,11 @@ def test_ssim_image():
     S2 = ssim(X, Y, win_size=11, gaussian_weights=True)
     assert(S1 < 0.3)
 
+    mssim0, S3 = ssim(X, Y, full=True)
+    assert_equal(S3.shape, X.shape)
+    mssim = ssim(X, Y)
+    assert_equal(mssim0, mssim)
+
 
 def test_ssim_multichannel():
     N = 100
@@ -55,9 +60,14 @@ def test_ssim_multichannel():
     S2 = ssim(Xc, Yc, win_size=3)
     assert_almost_equal(S1, S2)
 
+    # full case should return an image as well
+    m, S3 = ssim(Xc, Yc, full=True)
+    assert_equal(S3.shape, Xc.shape)
+
     # fail if win_size exceeds any non-channel dimension
     assert_raises(ValueError, ssim, Xc, Yc, win_size=7, multichannel=False)
 
+
 # NOTE: This test is known to randomly fail on some systems (Mac OS X 10.6)
 def test_ssim_grad():
     N = 30
@@ -71,6 +81,9 @@ def test_ssim_grad():
     assert g[0] < 0.05
     assert np.all(g[1] < 0.05)
 
+    mssim, grad, s = ssim(X, Y, dynamic_range=255, gradient=True, full=True)
+    assert np.all(grad < 0.05)
+
 
 def test_ssim_dtype():
     N = 30