From 82082ed4add783d672c05c462b0ad914d88a7bdd Mon Sep 17 00:00:00 2001 From: ckkelvinchan Date: Mon, 12 Apr 2021 09:18:17 +0800 Subject: [PATCH 1/3] Support Y-channel PSNR and SSIM --- mmedit/core/evaluation/metrics.py | 26 ++++++++++++++++++++++++-- tests/test_metrics.py | 12 ++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/mmedit/core/evaluation/metrics.py b/mmedit/core/evaluation/metrics.py index 16256ba52b..7ef60e4999 100644 --- 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', color_space=None): """Calculate PSNR (Peak Signal-to-Noise Ratio). Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio @@ -177,6 +177,8 @@ 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'. + color_space (str): The color space in which PSNR is computed. If None, + the images are not altered. Default: None. Returns: float: psnr result. @@ -191,6 +193,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(color_space, str) and color_space.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 color_space is not None: + raise ValueError(f'Wrong color space. 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 +246,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', color_space=None): """Calculate SSIM (structural similarity). Ref: @@ -255,6 +265,8 @@ 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'. + color_space (str): The color space in which SSIM is computed. If None, + the images are not altered. Default: None. Returns: float: ssim result. @@ -269,6 +281,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(color_space, str) and color_space.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 color_space is not None: + raise ValueError(f'Wrong color space. 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 index 29c3889e11..c1ea9ee697 100644 --- 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, color_space='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,9 @@ 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, color_space='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 +72,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 +89,9 @@ 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, color_space='Y') + np.testing.assert_almost_equal(ssim_result, 0.9987801) + def test_calculate_niqe(): img = mmcv.imread('tests/data/gt/baboon.png') From cf5a2db7b108d829bb5fcfe3edefda9c9467708c Mon Sep 17 00:00:00 2001 From: ckkelvinchan Date: Mon, 12 Apr 2021 09:55:29 +0800 Subject: [PATCH 2/3] fix bug --- mmedit/core/evaluation/metrics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mmedit/core/evaluation/metrics.py b/mmedit/core/evaluation/metrics.py index 7ef60e4999..1eb2eb180e 100644 --- a/mmedit/core/evaluation/metrics.py +++ b/mmedit/core/evaluation/metrics.py @@ -281,7 +281,7 @@ def ssim(img1, img2, crop_border=0, input_order='HWC', color_space=None): img1 = reorder_image(img1, input_order=input_order) img2 = reorder_image(img2, input_order=input_order) - if isinstance(color_space, str) and color_space.lower() == 'Y': + if isinstance(color_space, str) and color_space.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. From 1dd1009e9025b709c7e7d210f10a39375b4ca3f3 Mon Sep 17 00:00:00 2001 From: ckkelvinchan Date: Mon, 12 Apr 2021 10:49:54 +0800 Subject: [PATCH 3/3] update naming and unittest --- mmedit/core/evaluation/metrics.py | 30 +++++++++++++++++------------- tests/test_metrics.py | 10 +++++++--- 2 files changed, 24 insertions(+), 16 deletions(-) diff --git a/mmedit/core/evaluation/metrics.py b/mmedit/core/evaluation/metrics.py index 1eb2eb180e..5c85c42de7 100644 --- 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', color_space=None): +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,8 +177,10 @@ def psnr(img1, img2, crop_border=0, input_order='HWC', color_space=None): pixels are not involved in the PSNR calculation. Default: 0. input_order (str): Whether the input order is 'HWC' or 'CHW'. Default: 'HWC'. - color_space (str): The color space in which PSNR is computed. If None, - the images are not altered. Default: None. + 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. @@ -193,13 +195,13 @@ def psnr(img1, img2, crop_border=0, input_order='HWC', color_space=None): img1 = reorder_image(img1, input_order=input_order) img2 = reorder_image(img2, input_order=input_order) - if isinstance(color_space, str) and color_space.lower() == 'y': + 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 color_space is not None: - raise ValueError(f'Wrong color space. Supported values are ' - '"Y" and None') + 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] @@ -246,7 +248,7 @@ def _ssim(img1, img2): return ssim_map.mean() -def ssim(img1, img2, crop_border=0, input_order='HWC', color_space=None): +def ssim(img1, img2, crop_border=0, input_order='HWC', convert_to=None): """Calculate SSIM (structural similarity). Ref: @@ -265,8 +267,10 @@ def ssim(img1, img2, crop_border=0, input_order='HWC', color_space=None): pixels are not involved in the SSIM calculation. Default: 0. input_order (str): Whether the input order is 'HWC' or 'CHW'. Default: 'HWC'. - color_space (str): The color space in which SSIM is computed. If None, - the images are not altered. Default: None. + 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. @@ -281,14 +285,14 @@ def ssim(img1, img2, crop_border=0, input_order='HWC', color_space=None): img1 = reorder_image(img1, input_order=input_order) img2 = reorder_image(img2, input_order=input_order) - if isinstance(color_space, str) and color_space.lower() == 'y': + 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 color_space is not None: - raise ValueError(f'Wrong color space. Supported values are ' + elif convert_to is not None: + raise ValueError(f'Wrong color model. Supported values are ' '"Y" and None') if crop_border != 0: diff --git a/tests/test_metrics.py b/tests/test_metrics.py index c1ea9ee697..93f298f0ce 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -37,7 +37,7 @@ def test_calculate_psnr(): 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, color_space='ABC') + 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) @@ -53,7 +53,9 @@ 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, color_space='Y') + 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 @@ -89,7 +91,9 @@ 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, color_space='Y') + 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)