| 60 | |
| 61 | from scipy.ndimage import uniform_filter |
| 62 | def SSIM(im1, im2, mask, multichannel=True, heatmap=False, with_mask=True): |
| 63 | check_shape_equality(im1, im2) |
| 64 | |
| 65 | if multichannel: |
| 66 | # loop over channels |
| 67 | nch = im1.shape[-1] |
| 68 | mssim = np.empty(nch) |
| 69 | S = np.empty(im1.shape) |
| 70 | for ch in range(nch): |
| 71 | ch_result = SSIM(im1[..., ch], im2[..., ch], multichannel=False, mask=mask, heatmap=heatmap) |
| 72 | mssim[..., ch], S[...,ch] = ch_result |
| 73 | mssim = mssim.mean() |
| 74 | if heatmap: |
| 75 | S = np.mean(S, axis=2) |
| 76 | S = (S+1)/2 |
| 77 | S = (S * 255).astype(np.uint8) |
| 78 | error_heatmap = cv2.applyColorMap(S, cv2.COLORMAP_JET)*(1-mask) |
| 79 | return mssim, error_heatmap |
| 80 | else: |
| 81 | return mssim |
| 82 | |
| 83 | if im1.dtype != im2.dtype: |
| 84 | warn("Inputs have mismatched dtype. Setting data_range based on im1.dtype.", stacklevel=2) |
| 85 | dmin, dmax = dtype_range[im1.dtype.type] |
| 86 | data_range = dmax - dmin |
| 87 | |
| 88 | K1 = 0.01 |
| 89 | K2 = 0.03 |
| 90 | R = data_range |
| 91 | C1 = (K1 * R) ** 2 |
| 92 | C2 = (K2 * R) ** 2 |
| 93 | S = np.ones_like(im1) |
| 94 | win_size=3 |
| 95 | |
| 96 | # ndimage filters need floating point data |
| 97 | if not with_mask: |
| 98 | im1 = im1.astype(np.float64) |
| 99 | im2 = im2.astype(np.float64) |
| 100 | else: |
| 101 | mask = mask.squeeze() |
| 102 | im1 = (im1*(1-mask)).astype(np.float64) |
| 103 | im2 = (im2*(1-mask)).astype(np.float64) |
| 104 | |
| 105 | if heatmap: |
| 106 | ux = uniform_filter(im1, size=win_size) |
| 107 | uy = uniform_filter(im2, size=win_size) |
| 108 | |
| 109 | uxx = uniform_filter(im1*im1, size=win_size) |
| 110 | uyy = uniform_filter(im2*im2, size=win_size) |
| 111 | uxy = uniform_filter(im1*im2, size=win_size) |
| 112 | vx = uxx - ux*ux |
| 113 | vy = uyy - uy*uy |
| 114 | vxy = uxy - ux*uy |
| 115 | |
| 116 | A1, A2, B1, B2 = ((2 * ux * uy + C1, |
| 117 | 2 * vxy + C2, |
| 118 | ux ** 2 + uy ** 2 + C1, |
| 119 | vx + vy + C2)) |