from __future__ import division __all__ = ['structural_similarity'] import numpy as np from scipy.ndimage import uniform_filter, gaussian_filter from ..util.dtype import dtype_range from ..util.arraypad import crop def structural_similarity(X, Y, win_size=None, gradient=False, dynamic_range=None, multichannel=False, gaussian_weights=False, full=False, **kwargs): """Compute the mean structural similarity index between two images. Parameters ---------- X, Y : ndarray Image. Any dimensionality. win_size : int or None The side-length of the sliding window used in comparison. Must be an odd value. If `gaussian_weights` is True, this is ignored and the window size will depend on `sigma`. gradient : bool If True, also return the gradient. dynamic_range : int The dynamic range of the input image (distance between minimum and maximum possible values). By default, this is estimated from the image data-type. multichannel : int or None If True, treat the last dimension of the array as channels. Similarity calculations are done independently for each channel then averaged. gaussian_weights : bool If True, each patch 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 Other Parameters ---------------- use_sample_covariance : bool if True, normalize covariances by N-1 rather than, N where N is the number of pixels within the sliding window. K1 : float algorithm parameter, K1 (small constant, see [1]_) K2 : float algorithm parameter, K2 (small constant, see [1]_) sigma : float sigma for the Gaussian when `gaussian_weights` is True. Returns ------- mssim : float or ndarray The mean structural similarity over the image. grad : ndarray The gradient of the structural similarity index between X and Y [2]_. This is only returned if `gradient` is set to True. S : ndarray The full SSIM image. This is only returned if `full` is set to True. Notes ----- To match the implementation of Wang et. al. [1]_, set `gaussian_weights` to True, `sigma` to 1.5, and `use_sample_covariance` to False. References ---------- .. [1] Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). Image quality assessment: From error visibility to structural similarity. IEEE Transactions on Image Processing, 13, 600-612. https://ece.uwaterloo.ca/~z70wang/publications/ssim.pdf .. [2] Avanaki, A. N. (2009). Exact global histogram specification optimized for structural similarity. Optical Review, 16, 613-621. http://arxiv.org/abs/0901.0065 """ if not X.dtype == Y.dtype: raise ValueError('Input images must have the same dtype.') if not X.shape == Y.shape: raise ValueError('Input images must have the same dimensions.') if multichannel: # loop over channels args = dict(win_size=win_size, gradient=gradient, dynamic_range=dynamic_range, multichannel=False, gaussian_weights=gaussian_weights, full=full) args.update(kwargs) nch = X.shape[-1] mssim = np.empty(nch) if gradient: G = np.empty(X.shape) if full: S = np.empty(X.shape) for ch in range(nch): 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] = ch_result mssim = mssim.mean() if gradient and full: return mssim, G, S elif gradient: return mssim, G elif full: return mssim, S else: return mssim K1 = kwargs.pop('K1', 0.01) K2 = kwargs.pop('K2', 0.03) sigma = kwargs.pop('sigma', 1.5) if K1 < 0: raise ValueError("K1 must be positive") if K2 < 0: raise ValueError("K2 must be positive") if sigma < 0: raise ValueError("sigma must be positive") use_sample_covariance = kwargs.pop('use_sample_covariance', True) if win_size is None: if gaussian_weights: win_size = 11 # 11 to match Wang et. al. 2004 else: win_size = 7 # backwards compatibility if np.any((np.asarray(X.shape) - win_size) < 0): raise ValueError("win_size exceeds image extent") if not (win_size % 2 == 1): raise ValueError('Window size must be odd.') if dynamic_range is None: dmin, dmax = dtype_range[X.dtype.type] dynamic_range = dmax - dmin ndim = X.ndim if gaussian_weights: # sigma = 1.5 to approximately match filter in Wang et. al. 2004 # this ends up giving a 13-tap rather than 11-tap Gaussian filter_func = gaussian_filter filter_args = {'sigma': sigma} else: filter_func = uniform_filter filter_args = {'size': win_size} # ndimage filters need floating point data X = X.astype(np.float64) Y = Y.astype(np.float64) NP = win_size ** ndim # filter has already normalized by NP if use_sample_covariance: cov_norm = NP / (NP - 1) # sample covariance else: cov_norm = 1.0 # population covariance to match Wang et. al. 2004 # compute (weighted) means ux = filter_func(X, **filter_args) uy = filter_func(Y, **filter_args) # compute (weighted) variances and covariances uxx = filter_func(X * X, **filter_args) uyy = filter_func(Y * Y, **filter_args) uxy = filter_func(X * Y, **filter_args) vx = cov_norm * (uxx - ux * ux) vy = cov_norm * (uyy - uy * uy) vxy = cov_norm * (uxy - ux * uy) R = dynamic_range C1 = (K1 * R) ** 2 C2 = (K2 * R) ** 2 A1, A2, B1, B2 = ((2 * ux * uy + C1, 2 * vxy + C2, ux ** 2 + uy ** 2 + C1, vx + vy + C2)) D = B1 * B2 S = (A1 * A2) / D # to avoid edge effects will ignore filter radius strip around edges pad = (win_size - 1) // 2 # compute (weighted) mean of ssim mssim = crop(S, pad).mean() if gradient: # The following is Eqs. 7-8 of Avanaki 2009. grad = filter_func(A1 / D, **filter_args) * X grad += filter_func(-S / B2, **filter_args) * Y grad += filter_func((ux * (A2 - A1) - uy * (B2 - B1) * S) / D, **filter_args) grad *= (2 / X.size) if full: return mssim, grad, S else: return mssim, grad else: if full: return mssim, S else: return mssim