| 33 | |
| 34 | |
| 35 | def ssim(img1, img2): |
| 36 | assert img1.shape == img2.shape |
| 37 | assert img1.ndim == 2 and img2.ndim == 2 |
| 38 | C1 = (0.01 * 255)**2 |
| 39 | C2 = (0.03 * 255)**2 |
| 40 | |
| 41 | img1 = img1.astype(np.float64) |
| 42 | img2 = img2.astype(np.float64) |
| 43 | kernel = cv2.getGaussianKernel(11, 1.5) |
| 44 | window = np.outer(kernel, kernel.transpose()) |
| 45 | |
| 46 | mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5] # valid |
| 47 | mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5] |
| 48 | mu1_sq = mu1**2 |
| 49 | mu2_sq = mu2**2 |
| 50 | mu1_mu2 = mu1 * mu2 |
| 51 | sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq |
| 52 | sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq |
| 53 | sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2 |
| 54 | |
| 55 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * |
| 56 | (sigma1_sq + sigma2_sq + C2)) |
| 57 | return ssim_map.mean() |
| 58 | |
| 59 | |
| 60 | def calculate_psnr(img1, img2): |