calculate SSIM the same outputs as MATLAB's img1, img2: [0, 255]
(img1, img2, border=0)
| 640 | # SSIM |
| 641 | # -------------------------------------------- |
| 642 | def calculate_ssim(img1, img2, border=0): |
| 643 | '''calculate SSIM |
| 644 | the same outputs as MATLAB's |
| 645 | img1, img2: [0, 255] |
| 646 | ''' |
| 647 | #img1 = img1.squeeze() |
| 648 | #img2 = img2.squeeze() |
| 649 | if not img1.shape == img2.shape: |
| 650 | raise ValueError('Input images must have the same dimensions.') |
| 651 | h, w = img1.shape[:2] |
| 652 | img1 = img1[border:h-border, border:w-border] |
| 653 | img2 = img2[border:h-border, border:w-border] |
| 654 | |
| 655 | if img1.ndim == 2: |
| 656 | return ssim(img1, img2) |
| 657 | elif img1.ndim == 3: |
| 658 | if img1.shape[2] == 3: |
| 659 | ssims = [] |
| 660 | for i in range(3): |
| 661 | ssims.append(ssim(img1[:,:,i], img2[:,:,i])) |
| 662 | return np.array(ssims).mean() |
| 663 | elif img1.shape[2] == 1: |
| 664 | return ssim(np.squeeze(img1), np.squeeze(img2)) |
| 665 | else: |
| 666 | raise ValueError('Wrong input image dimensions.') |
| 667 | |
| 668 | |
| 669 | def ssim(img1, img2): |