Merge 1dd1009 into 8725d92

open-mmlab · Apr 12, 2021 · fbaaef0 · fbaaef0
2 parents 8725d92 + 1dd1009
commit fbaaef0
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 2 deletions.
diff --git a/mmedit/core/evaluation/metrics.py b/mmedit/core/evaluation/metrics.py
@@ -165,7 +165,7 @@ def reorder_image(img, input_order='HWC'):
     return img
 
 
-def psnr(img1, img2, crop_border=0, input_order='HWC'):
+def psnr(img1, img2, crop_border=0, input_order='HWC', convert_to=None):
     """Calculate PSNR (Peak Signal-to-Noise Ratio).
 
     Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
@@ -177,6 +177,10 @@ def psnr(img1, img2, crop_border=0, input_order='HWC'):
             pixels are not involved in the PSNR calculation. Default: 0.
         input_order (str): Whether the input order is 'HWC' or 'CHW'.
             Default: 'HWC'.
+        convert_to (str): Whether to convert the images to other color models.
+            If None, the images are not altered. When computing for 'Y',
+            the images are assumed to be in BGR order. Options are 'Y' and
+            None. Default: None.
 
     Returns:
         float: psnr result.
@@ -191,6 +195,14 @@ def psnr(img1, img2, crop_border=0, input_order='HWC'):
     img1 = reorder_image(img1, input_order=input_order)
     img2 = reorder_image(img2, input_order=input_order)
 
+    if isinstance(convert_to, str) and convert_to.lower() == 'y':
+        img1, img2 = img1.astype(np.float32), img2.astype(np.float32)
+        img1 = mmcv.bgr2ycbcr(img1 / 255., y_only=True) * 255.
+        img2 = mmcv.bgr2ycbcr(img2 / 255., y_only=True) * 255.
+    elif convert_to is not None:
+        raise ValueError(f'Wrong color model. Supported values are '
+                         '"Y" and None.')
+
     if crop_border != 0:
         img1 = img1[crop_border:-crop_border, crop_border:-crop_border, None]
         img2 = img2[crop_border:-crop_border, crop_border:-crop_border, None]
@@ -236,7 +248,7 @@ def _ssim(img1, img2):
     return ssim_map.mean()
 
 
-def ssim(img1, img2, crop_border=0, input_order='HWC'):
+def ssim(img1, img2, crop_border=0, input_order='HWC', convert_to=None):
     """Calculate SSIM (structural similarity).
 
     Ref:
@@ -255,6 +267,10 @@ def ssim(img1, img2, crop_border=0, input_order='HWC'):
             pixels are not involved in the SSIM calculation. Default: 0.
         input_order (str): Whether the input order is 'HWC' or 'CHW'.
             Default: 'HWC'.
+        convert_to (str): Whether to convert the images to other color models.
+            If None, the images are not altered. When computing for 'Y',
+            the images are assumed to be in BGR order. Options are 'Y' and
+            None. Default: None.
 
     Returns:
         float: ssim result.
@@ -269,6 +285,16 @@ def ssim(img1, img2, crop_border=0, input_order='HWC'):
     img1 = reorder_image(img1, input_order=input_order)
     img2 = reorder_image(img2, input_order=input_order)
 
+    if isinstance(convert_to, str) and convert_to.lower() == 'y':
+        img1, img2 = img1.astype(np.float32), img2.astype(np.float32)
+        img1 = mmcv.bgr2ycbcr(img1 / 255., y_only=True) * 255.
+        img2 = mmcv.bgr2ycbcr(img2 / 255., y_only=True) * 255.
+        img1 = np.expand_dims(img1, axis=2)
+        img2 = np.expand_dims(img2, axis=2)
+    elif convert_to is not None:
+        raise ValueError(f'Wrong color model. Supported values are '
+                         '"Y" and None')
+
     if crop_border != 0:
         img1 = img1[crop_border:-crop_border, crop_border:-crop_border, None]
         img2 = img2[crop_border:-crop_border, crop_border:-crop_border, None]

diff --git a/tests/test_metrics.py b/tests/test_metrics.py
@@ -36,6 +36,9 @@ def test_calculate_psnr():
     with pytest.raises(ValueError):
         psnr(img_hw_1, img_hw_2, crop_border=0, input_order='HH')
 
+    with pytest.raises(ValueError):
+        psnr(img_hw_1, img_hw_2, crop_border=0, convert_to='ABC')
+
     psnr_result = psnr(img_hw_1, img_hw_2, crop_border=0)
     np.testing.assert_almost_equal(psnr_result, 48.1308036)
     psnr_result = psnr(img_hwc_1, img_hwc_2, crop_border=0, input_order='HWC')
@@ -50,6 +53,11 @@ def test_calculate_psnr():
     psnr_result = psnr(img_chw_1, img_chw_2, crop_border=4, input_order='CHW')
     np.testing.assert_almost_equal(psnr_result, 48.1308036)
 
+    psnr_result = psnr(img_hwc_1, img_hwc_2, crop_border=0, convert_to=None)
+    np.testing.assert_almost_equal(psnr_result, 48.1308036)
+    psnr_result = psnr(img_hwc_1, img_hwc_2, crop_border=0, convert_to='Y')
+    np.testing.assert_almost_equal(psnr_result, 49.4527218)
+
     # test float inf
     psnr_result = psnr(img_hw_1, img_hw_1, crop_border=0)
     assert psnr_result == float('inf')
@@ -66,6 +74,9 @@ def test_calculate_ssim():
     with pytest.raises(ValueError):
         ssim(img_hw_1, img_hw_2, crop_border=0, input_order='HH')
 
+    with pytest.raises(ValueError):
+        ssim(img_hw_1, img_hw_2, crop_border=0, input_order='ABC')
+
     ssim_result = ssim(img_hw_1, img_hw_2, crop_border=0)
     np.testing.assert_almost_equal(ssim_result, 0.9130623)
     ssim_result = ssim(img_hwc_1, img_hwc_2, crop_border=0, input_order='HWC')
@@ -80,6 +91,11 @@ def test_calculate_ssim():
     ssim_result = ssim(img_chw_1, img_chw_2, crop_border=4, input_order='CHW')
     np.testing.assert_almost_equal(ssim_result, 0.9130623)
 
+    ssim_result = ssim(img_hwc_1, img_hwc_2, crop_border=0, convert_to=None)
+    np.testing.assert_almost_equal(ssim_result, 0.9130623)
+    ssim_result = ssim(img_hwc_1, img_hwc_2, crop_border=0, convert_to='Y')
+    np.testing.assert_almost_equal(ssim_result, 0.9987801)
+
 
 def test_calculate_niqe():
     img = mmcv.imread('tests/data/gt/baboon.png')