| 15 | import os |
| 16 | |
| 17 | def PSNR(image_true, image_test, mask, heatmap=False, with_mask=True): |
| 18 | check_shape_equality(image_true, image_test) |
| 19 | |
| 20 | if image_true.dtype != image_test.dtype: |
| 21 | warn("Inputs have mismatched dtype. Setting data_range based on " |
| 22 | "im_true.", stacklevel=2) |
| 23 | dmin, dmax = dtype_range[image_true.dtype.type] |
| 24 | true_min, true_max = np.min(image_true), np.max(image_true) |
| 25 | if true_max > dmax or true_min < dmin: |
| 26 | raise ValueError( |
| 27 | "im_true has intensity values outside the range expected for " |
| 28 | "its data type. Please manually specify the data_range") |
| 29 | if true_min >= 0: |
| 30 | # most common case (255 for uint8, 1 for float) |
| 31 | data_range = dmax |
| 32 | else: |
| 33 | data_range = dmax - dmin |
| 34 | |
| 35 | image_true = image_true.astype(np.float64) |
| 36 | image_test = image_test.astype(np.float64) |
| 37 | if not with_mask: |
| 38 | error_mask = ((image_true - image_test) ** 2).astype(np.float64) |
| 39 | err = np.mean(error_mask, dtype=np.float64) |
| 40 | else: |
| 41 | cnt = np.count_nonzero(1-mask) * image_true.shape[2] |
| 42 | error_mask = ((image_true*(1-mask) - image_test*(1-mask))**2).astype(np.float64) |
| 43 | sum = np.sum(error_mask) |
| 44 | err = sum/cnt |
| 45 | |
| 46 | score = 10 * np.log10((data_range ** 2) / err) |
| 47 | if heatmap: |
| 48 | error_mask = np.mean(error_mask, axis=2) |
| 49 | # np.save('./error_mask.npy', error_mask) |
| 50 | |
| 51 | x,y = np.nonzero(error_mask) |
| 52 | error_mask[x,y] = 10 * np.log10((data_range ** 2) / error_mask[x,y]) |
| 53 | error_mask = error_mask / 30 |
| 54 | |
| 55 | error_mask = (error_mask * 255).astype(np.uint8) |
| 56 | error_heatmap = cv2.applyColorMap(error_mask, cv2.COLORMAP_JET)*(1-mask) |
| 57 | return score, error_heatmap |
| 58 | else: |
| 59 | return score |
| 60 | |
| 61 | from scipy.ndimage import uniform_filter |
| 62 | def SSIM(im1, im2, mask, multichannel=True, heatmap=False, with_mask=True): |