diff --git a/mmedit/core/evaluation/metrics.py b/mmedit/core/evaluation/metrics.py index 16256ba52b..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'): +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 index 29c3889e11..93f298f0ce 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, 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')