diff --git a/skimage/measure/_ssim.py b/skimage/measure/_ssim.py index a93571e8..f0d8357c 100644 --- a/skimage/measure/_ssim.py +++ b/skimage/measure/_ssim.py @@ -5,7 +5,7 @@ __all__ = ['ssim'] import numpy as np from numpy.lib import stride_tricks -def _as_windows(X, win_size=7): +def _as_windows(X, win_size=7, flatten_first_axis=True): """Re-stride an array to simulate a sliding window. Parameters @@ -15,7 +15,7 @@ def _as_windows(X, win_size=7): Returns ------- - window : (N, win_size, win_size) ndarray + window : (N, M, win_size, win_size) ndarray Sliding windows. """ @@ -35,12 +35,11 @@ def _as_windows(X, win_size=7): new_shape = (new_rows, new_cols, win_size, win_size) windows = stride_tricks.as_strided(X, shape=new_shape, strides=new_strides) - windows = windows.reshape((-1, win_size, win_size)) return windows -def ssim(X, Y, win_size=7, dynamic_range=255): +def ssim(X, Y, win_size=7, gradient=False, dynamic_range=255): """Compute the structural similarity index between two images. Parameters @@ -54,11 +53,16 @@ def ssim(X, Y, win_size=7, dynamic_range=255): Dynamic range of the input image (distance between minimum and maximum possible values). This should eventually be auto-computed, but just specifying it manually for now. + gradient : bool + If True, also return the gradient. Returns ------- s : float Strucutural similarity. + grad : (N * N,) ndarray + Gradient of the structural similarity index between X and Y. + This is only returned if `gradient` is set to True. References ---------- @@ -72,33 +76,27 @@ def ssim(X, Y, win_size=7, dynamic_range=255): raise ValueError('Input images must have the same dtype.') if not X.shape == Y.shape: - raise ValueError('Inout images must have the same dimensions.') + raise ValueError('Input images must have the same dimensions.') - import time - - tic = time.time() + if not (win_size % 2 == 1): + raise ValueError('Window size must be odd.') XW = _as_windows(X, win_size=win_size) YW = _as_windows(Y, win_size=win_size) - tic = time.time() - - # Flatten windows - XW = XW.reshape(XW.shape[0], -1) - YW = YW.reshape(YW.shape[0], -1) + NS = len(XW) + NP = win_size * win_size - ux = np.mean(XW, axis=1) - uy = np.mean(YW, axis=1) - - tic = time.time() + ux = np.mean(np.mean(XW, axis=2), axis=2) + uy = np.mean(np.mean(YW, axis=2), axis=2) # Compute variances var(X), var(Y) and var(X, Y) cov_norm = 1 / (win_size**2 - 1) - XWM = XW - ux[:, None] - YWM = YW - uy[:, None] - vx = cov_norm * np.sum(XWM**2, axis=1) - vy = cov_norm * np.sum(YWM**2, axis=1) - vxy = cov_norm * np.sum(XWM * YWM, axis=1) + XWM = XW - ux[..., None, None] + YWM = YW - uy[..., None, None] + vx = cov_norm * np.sum(np.sum(XWM**2, axis=2), axis=2) + vy = cov_norm * np.sum(np.sum(YWM**2, axis=2), axis=2) + vxy = cov_norm * np.sum(np.sum(XWM * YWM, axis=2), axis=2) R = dynamic_range K1 = 0.01 @@ -106,5 +104,29 @@ def ssim(X, Y, win_size=7, dynamic_range=255): C1 = (K1 * R)**2 C2 = (K2 * R)**2 - return np.mean(((2 * ux * uy + C1) * (2 * vxy + C2)) / \ - ((ux**2 + uy**2 + C1) * (vx + vy + C2))) + A1, A2, B1, B2 = (v[..., None, None] for v in + (2 * ux * uy + C1, + 2 * vxy + C2, + ux**2 + uy**2 + C1, + vx + vy + C2)) + + S = np.mean((A1 * A2) / (B1 * B2)) + + if gradient: + local_grad = 2 / (NP * B1**2 * B2**2) * \ + ( + A1 * B1 * (B2 * XW - A2 * YW) - \ + B1 * B2 * (A2 - A1) * ux[..., None, None] + \ + A1 * A2 * (B1 - B2) * uy[..., None, None] + ) + + grad = np.zeros_like(X, dtype=float) + OW = _as_windows(grad, win_size=win_size) + + OW += local_grad + grad /= NS + + return S, grad + + else: + return S diff --git a/skimage/measure/tests/test_ssim.py b/skimage/measure/tests/test_ssim.py index 8e6684b0..f140a48d 100644 --- a/skimage/measure/tests/test_ssim.py +++ b/skimage/measure/tests/test_ssim.py @@ -2,6 +2,7 @@ import numpy as np from numpy.testing import assert_equal from skimage.measure._ssim import ssim, _as_windows +import scipy.optimize as opt def test_ssim_patch_range(): N = 51 @@ -14,12 +15,12 @@ def test_ssim_patch_range(): def test_as_windows(): X = np.arange(100).reshape((10, 10)) W = _as_windows(X, win_size=7) - assert_equal(len(W), 16) + assert_equal(W.shape[:2], (4, 4)) W = _as_windows(X, win_size=3) - assert_equal(W[0], [[0, 1, 2], - [10, 11, 12], - [20, 21, 22]]) + assert_equal(W[0, 0], [[0, 1, 2], + [10, 11, 12], + [20, 21, 22]]) def test_ssim_image(): N = 100 @@ -32,5 +33,25 @@ def test_ssim_image(): S1 = ssim(X, Y, win_size=3) assert(S1 < 0.3) +def test_ssim_grad(): + N = 30 + X = np.random.random((N, N)) + Y = np.random.random((N, N)) + + def func(Y): + return ssim(X, Y) + + def grad(Y): + return ssim(X, Y, gradient=True)[1] + + assert(np.all(opt.check_grad(func, grad, Y) < 0.05)) + +# N = 200 +# X = np.random.random((N, N)) +# Y = np.random.random((N, N)) + +# assert(np.all(np.abs(ssim(X, Y, gradient=True))[1] < 1e-2)) + + if __name__ == "__main__": np.testing.run_module_suite()