def main():
"""Calculate PSNR and SSIM for images.
Configurations:
folder_gt (str): Path to gt (Ground-Truth).
folder_restored (str): Path to restored images.
crop_border (int): Crop border for each side.
suffix (str): Suffix for restored images.
test_y_channel (bool): If True, test Y channel (In MatLab YCbCr format)
If False, test RGB channels.
"""
# Configurations
# -------------------------------------------------------------------------
folder_gt = 'datasets/val_set14/Set14'
folder_restored = 'results/exp/visualization/val_set14'
crop_border = 4
suffix = '_expname'
test_y_channel = False
# -------------------------------------------------------------------------
psnr_all = []
ssim_all = []
img_list = sorted(mmcv.scandir(folder_gt, recursive=True))
if test_y_channel:
print('Testing Y channel.')
else:
print('Testing RGB channels.')
for i, img_path in enumerate(img_list):
basename, ext = osp.splitext(osp.basename(img_path))
img_gt = mmcv.imread(
osp.join(folder_gt, img_path), flag='unchanged').astype(
np.float32) / 255.
img_restored = mmcv.imread(
osp.join(folder_restored, basename + suffix + ext),
flag='unchanged').astype(np.float32) / 255.
if test_y_channel and img_gt.ndim == 3 and img_gt.shape[2] == 3:
img_gt = mmcv.bgr2ycbcr(img_gt, y_only=True)
img_restored = mmcv.bgr2ycbcr(img_restored, y_only=True)
# calculate PSNR and SSIM
psnr = calculate_psnr(
img_gt * 255,
img_restored * 255,
crop_border=crop_border,
input_order='HWC')
ssim = calculate_ssim(
img_gt * 255,
img_restored * 255,
crop_border=crop_border,
input_order='HWC')
print(f'{i+1:3d}: {basename:25}. \tPSNR: {psnr:.6f} dB, '
f'\tSSIM: {ssim:.6f}')
psnr_all.append(psnr)
ssim_all.append(ssim)
print(f'Average: PSNR: {sum(psnr_all) / len(psnr_all):.6f} dB, '
f'SSIM: {sum(ssim_all) / len(ssim_all):.6f}')
def ssim(img1, img2, crop_border=0, input_order='HWC', convert_to=None):
"""Calculate SSIM (structural similarity).
Ref:
Image quality assessment: From error visibility to structural similarity
The results are the same as that of the official released MATLAB code in
https://ece.uwaterloo.ca/~z70wang/research/ssim/.
For three-channel images, SSIM is calculated for each channel and then
averaged.
Args:
img1 (ndarray): Images with range [0, 255].
img2 (ndarray): Images with range [0, 255].
crop_border (int): Cropped pixels in each edges of an image. These
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.
"""
assert img1.shape == img2.shape, (
f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
if input_order not in ['HWC', 'CHW']:
raise ValueError(
f'Wrong input_order {input_order}. Supported input_orders are '
'"HWC" and "CHW"')
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]
ssims = []
for i in range(img1.shape[2]):
ssims.append(_ssim(img1[..., i], img2[..., i]))
return np.array(ssims).mean()
def test_bgr2ycbcr():
# float32
in_img = np.random.rand(10, 10, 3).astype(np.float32)
out_img = mmcv.bgr2ycbcr(in_img)
computed_ycbcr = np.empty_like(in_img)
for i in range(in_img.shape[0]):
for j in range(in_img.shape[1]):
b, g, r = in_img[i, j]
y = 16 + r * 65.481 + g * 128.553 + b * 24.966
cb = 128 - r * 37.797 - g * 74.203 + b * 112.0
cr = 128 + r * 112.0 - g * 93.786 - b * 18.214
computed_ycbcr[i, j, :] = [y, cb, cr]
computed_ycbcr /= 255.
assert_array_almost_equal(out_img, computed_ycbcr, decimal=2)
# y_only=True
in_img = np.random.rand(10, 10, 3).astype(np.float32)
out_img = mmcv.bgr2ycbcr(in_img, y_only=True)
computed_y = np.empty_like(out_img, dtype=out_img.dtype)
for i in range(in_img.shape[0]):
for j in range(in_img.shape[1]):
b, g, r = in_img[i, j]
y = 16 + r * 65.481 + g * 128.553 + b * 24.966
computed_y[i, j] = y
computed_y /= 255.
assert_array_almost_equal(out_img, computed_y, decimal=2)
# uint8
in_img = (np.random.rand(10, 10, 3) * 255).astype(np.uint8)
out_img = mmcv.bgr2ycbcr(in_img)
computed_ycbcr = np.empty_like(in_img)
in_img = in_img / 255.
for i in range(in_img.shape[0]):
for j in range(in_img.shape[1]):
b, g, r = in_img[i, j]
y = 16 + r * 65.481 + g * 128.553 + b * 24.966
cb = 128 - r * 37.797 - g * 74.203 + b * 112.0
cr = 128 + r * 112.0 - g * 93.786 - b * 18.214
y, cb, cr = y.round(), cb.round(), cr.round()
computed_ycbcr[i, j, :] = [y, cb, cr]
assert_image_almost_equal(out_img, computed_ycbcr)
# y_only = True
in_img = (np.random.rand(10, 10, 3) * 255).astype(np.uint8)
out_img = mmcv.bgr2ycbcr(in_img, y_only=True)
computed_y = np.empty_like(out_img, dtype=out_img.dtype)
in_img = in_img / 255.
for i in range(in_img.shape[0]):
for j in range(in_img.shape[1]):
b, g, r = in_img[i, j]
y = 16 + r * 65.481 + g * 128.553 + b * 24.966
y = y.round()
computed_y[i, j] = y
assert_image_almost_equal(out_img, computed_y)
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
Args:
img1 (ndarray): Images with range [0, 255].
img2 (ndarray): Images with range [0, 255].
crop_border (int): Cropped pixels in each edges of an image. These
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.
"""
assert img1.shape == img2.shape, (
f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
if input_order not in ['HWC', 'CHW']:
raise ValueError(
f'Wrong input_order {input_order}. Supported input_orders are '
'"HWC" and "CHW"')
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]
mse_value = np.mean((img1 - img2)**2)
if mse_value == 0:
return float('inf')
return 20. * np.log10(255. / np.sqrt(mse_value))
def __call__(self, results):
"""Call function.
Args:
results (dict): A dict containing the necessary information and
data for augmentation.
Returns:
dict: A dict containing the processed data and information.
"""
if self.file_client is None:
self.file_client = FileClient(self.io_backend, **self.kwargs)
filepaths = results[f'{self.key}_path']
if not isinstance(filepaths, list):
raise TypeError(
f'filepath should be list, but got {type(filepaths)}')
filepaths = [str(v) for v in filepaths]
imgs = []
shapes = []
if self.save_original_img:
ori_imgs = []
for filepath in filepaths:
img_bytes = self.file_client.get(filepath)
img = mmcv.imfrombytes(img_bytes,
flag=self.flag,
channel_order=self.channel_order) # HWC
# convert to y-channel, if specified
if self.convert_to is not None:
if self.channel_order == 'bgr' and self.convert_to.lower(
) == 'y':
img = mmcv.bgr2ycbcr(img, y_only=True)
elif self.channel_order == 'rgb':
img = mmcv.rgb2ycbcr(img, y_only=True)
else:
raise ValueError('Currently support only "bgr2ycbcr" or '
'"bgr2ycbcr".')
if img.ndim == 2:
img = np.expand_dims(img, axis=2)
imgs.append(img)
shapes.append(img.shape)
if self.save_original_img:
ori_imgs.append(img.copy())
results[self.key] = imgs
results[f'{self.key}_path'] = filepaths
results[f'{self.key}_ori_shape'] = shapes
if self.save_original_img:
results[f'ori_{self.key}'] = ori_imgs
return results
def niqe(img, crop_border, input_order='HWC', convert_to='y'):
"""Calculate NIQE (Natural Image Quality Evaluator) metric.
Ref: Making a "Completely Blind" Image Quality Analyzer.
This implementation could produce almost the same results as the official
MATLAB codes: http://live.ece.utexas.edu/research/quality/niqe_release.zip
We use the official params estimated from the pristine dataset.
We use the recommended block size (96, 96) without overlaps.
Args:
img (ndarray): Input image whose quality needs to be computed.
The input image must be in range [0, 255] with float/int type.
The input_order of image can be 'HW' or 'HWC' or 'CHW'. (BGR order)
If the input order is 'HWC' or 'CHW', it will be converted to gray
or Y (of YCbCr) image according to the ``convert_to`` argument.
crop_border (int): Cropped pixels in each edge of an image. These
pixels are not involved in the metric calculation.
input_order (str): Whether the input order is 'HW', 'HWC' or 'CHW'.
Default: 'HWC'.
convert_to (str): Whether converted to 'y' (of MATLAB YCbCr) or 'gray'.
Default: 'y'.
Returns:
float: NIQE result.
"""
# we use the official params estimated from the pristine dataset.
niqe_pris_params = np.load('mmedit/core/evaluation/niqe_pris_params.npz')
mu_pris_param = niqe_pris_params['mu_pris_param']
cov_pris_param = niqe_pris_params['cov_pris_param']
gaussian_window = niqe_pris_params['gaussian_window']
img = img.astype(np.float32)
if input_order != 'HW':
img = reorder_image(img, input_order=input_order)
if convert_to == 'y':
img = mmcv.bgr2ycbcr(img / 255., y_only=True) * 255.
elif convert_to == 'gray':
img = mmcv.bgr2gray(img / 255., cv2.COLOR_BGR2GRAY) * 255.
img = np.squeeze(img)
if crop_border != 0:
img = img[crop_border:-crop_border, crop_border:-crop_border]
# round to follow official implementation
img = img.round()
niqe_result = niqe_core(img, mu_pris_param, cov_pris_param,
gaussian_window)
return niqe_result
def __call__(self, results):
"""Call function.
Args:
results (dict): A dict containing the necessary information and
data for augmentation.
Returns:
dict: A dict containing the processed data and information.
"""
filepath = str(results[f'{self.key}_path'])
if self.file_client is None:
self.file_client = FileClient(self.io_backend, **self.kwargs)
if self.use_cache:
if self.cache is None:
self.cache = dict()
if filepath in self.cache:
img = self.cache[filepath]
else:
img_bytes = self.file_client.get(filepath)
img = mmcv.imfrombytes(img_bytes,
flag=self.flag,
channel_order=self.channel_order,
backend=self.backend) # HWC
self.cache[filepath] = img
else:
img_bytes = self.file_client.get(filepath)
img = mmcv.imfrombytes(img_bytes,
flag=self.flag,
channel_order=self.channel_order,
backend=self.backend) # HWC
if self.convert_to is not None:
if self.channel_order == 'bgr' and self.convert_to.lower() == 'y':
img = mmcv.bgr2ycbcr(img, y_only=True)
elif self.channel_order == 'rgb':
img = mmcv.rgb2ycbcr(img, y_only=True)
else:
raise ValueError('Currently support only "bgr2ycbcr" or '
'"bgr2ycbcr".')
if img.ndim == 2:
img = np.expand_dims(img, axis=2)
results[self.key] = img
results[f'{self.key}_path'] = filepath
results[f'{self.key}_ori_shape'] = img.shape
if self.save_original_img:
results[f'ori_{self.key}'] = img.copy()
return results
def to_y_channel(img):
"""Change to Y channel of YCbCr.
Args:
img (ndarray): Images with range [0, 255].
Returns:
(ndarray): Images with range [0, 255] (float type) without round.
"""
img = img.astype(np.float32) / 255.
if img.ndim == 3 and img.shape[2] == 3:
img = mmcv.bgr2ycbcr(img, y_only=True)
img = img[..., None]
return img * 255.
def test_load_image_from_file():
path_baboon = Path(
__file__).parent.parent.parent / 'data' / 'gt' / 'baboon.png'
img_baboon = mmcv.imread(str(path_baboon), flag='color')
path_baboon_x4 = Path(
__file__).parent.parent.parent / 'data' / 'lq' / 'baboon_x4.png'
img_baboon_x4 = mmcv.imread(str(path_baboon_x4), flag='color')
# read gt image
# input path is Path object
results = dict(gt_path=path_baboon)
config = dict(io_backend='disk', key='gt')
image_loader = LoadImageFromFile(**config)
results = image_loader(results)
assert results['gt'].shape == (480, 500, 3)
np.testing.assert_almost_equal(results['gt'], img_baboon)
assert results['gt_path'] == str(path_baboon)
# input path is str
results = dict(gt_path=str(path_baboon))
results = image_loader(results)
assert results['gt'].shape == (480, 500, 3)
np.testing.assert_almost_equal(results['gt'], img_baboon)
assert results['gt_path'] == str(path_baboon)
# read lq image
# input path is Path object
results = dict(lq_path=path_baboon_x4)
config = dict(io_backend='disk', key='lq')
image_loader = LoadImageFromFile(**config)
results = image_loader(results)
assert results['lq'].shape == (120, 125, 3)
np.testing.assert_almost_equal(results['lq'], img_baboon_x4)
assert results['lq_path'] == str(path_baboon_x4)
# input path is str
results = dict(lq_path=str(path_baboon_x4))
results = image_loader(results)
assert results['lq'].shape == (120, 125, 3)
np.testing.assert_almost_equal(results['lq'], img_baboon_x4)
assert results['lq_path'] == str(path_baboon_x4)
assert repr(image_loader) == (
image_loader.__class__.__name__ +
('(io_backend=disk, key=lq, '
'flag=color, save_original_img=False, channel_order=bgr, '
'use_cache=False)'))
results = dict(lq_path=path_baboon_x4)
config = dict(io_backend='disk',
key='lq',
flag='grayscale',
save_original_img=True)
image_loader = LoadImageFromFile(**config)
results = image_loader(results)
assert results['lq'].shape == (120, 125)
assert results['lq_ori_shape'] == (120, 125)
np.testing.assert_almost_equal(results['ori_lq'], results['lq'])
assert id(results['ori_lq']) != id(results['lq'])
# test: use_cache
results = dict(gt_path=path_baboon)
config = dict(io_backend='disk', key='gt', use_cache=True)
image_loader = LoadImageFromFile(**config)
assert image_loader.cache is None
assert repr(image_loader) == (
image_loader.__class__.__name__ +
('(io_backend=disk, key=gt, '
'flag=color, save_original_img=False, channel_order=bgr, '
'use_cache=True)'))
results = image_loader(results)
assert image_loader.cache is not None
assert str(path_baboon) in image_loader.cache
assert results['gt'].shape == (480, 500, 3)
assert results['gt_path'] == str(path_baboon)
np.testing.assert_almost_equal(results['gt'], img_baboon)
# convert to y-channel (bgr2y)
results = dict(gt_path=path_baboon)
config = dict(io_backend='disk', key='gt', convert_to='y')
image_loader = LoadImageFromFile(**config)
results = image_loader(results)
assert results['gt'].shape == (480, 500, 1)
img_baboon_y = mmcv.bgr2ycbcr(img_baboon, y_only=True)
img_baboon_y = np.expand_dims(img_baboon_y, axis=2)
np.testing.assert_almost_equal(results['gt'], img_baboon_y)
assert results['gt_path'] == str(path_baboon)
# convert to y-channel (rgb2y)
results = dict(gt_path=path_baboon)
config = dict(io_backend='disk',
key='gt',
channel_order='rgb',
convert_to='y')
image_loader = LoadImageFromFile(**config)
results = image_loader(results)
assert results['gt'].shape == (480, 500, 1)
img_baboon_y = mmcv.bgr2ycbcr(img_baboon, y_only=True)
img_baboon_y = np.expand_dims(img_baboon_y, axis=2)
np.testing.assert_almost_equal(results['gt'], img_baboon_y)
assert results['gt_path'] == str(path_baboon)
# convert to y-channel (ValueError)
results = dict(gt_path=path_baboon)
config = dict(io_backend='disk', key='gt', convert_to='abc')
image_loader = LoadImageFromFile(**config)
with pytest.raises(ValueError):
results = image_loader(results)
def test_load_image_from_file_list():
path_baboon = Path(
__file__).parent.parent.parent / 'data' / 'gt' / 'baboon.png'
img_baboon = mmcv.imread(str(path_baboon), flag='color')
path_baboon_x4 = Path(
__file__).parent.parent.parent / 'data' / 'lq' / 'baboon_x4.png'
img_baboon_x4 = mmcv.imread(str(path_baboon_x4), flag='color')
# input path is Path object
results = dict(lq_path=[path_baboon_x4, path_baboon])
config = dict(io_backend='disk', key='lq')
image_loader = LoadImageFromFileList(**config)
results = image_loader(results)
np.testing.assert_almost_equal(results['lq'][0], img_baboon_x4)
np.testing.assert_almost_equal(results['lq'][1], img_baboon)
assert results['lq_ori_shape'] == [(120, 125, 3), (480, 500, 3)]
assert results['lq_path'] == [str(path_baboon_x4), str(path_baboon)]
# input path is str
results = dict(lq_path=[str(path_baboon_x4), str(path_baboon)])
config = dict(io_backend='disk', key='lq')
image_loader = LoadImageFromFileList(**config)
results = image_loader(results)
np.testing.assert_almost_equal(results['lq'][0], img_baboon_x4)
np.testing.assert_almost_equal(results['lq'][1], img_baboon)
assert results['lq_path'] == [str(path_baboon_x4), str(path_baboon)]
# save ori_img
results = dict(lq_path=[path_baboon_x4])
config = dict(io_backend='disk', key='lq', save_original_img=True)
image_loader = LoadImageFromFileList(**config)
results = image_loader(results)
np.testing.assert_almost_equal(results['lq'][0], img_baboon_x4)
assert results['lq_ori_shape'] == [(120, 125, 3)]
assert results['lq_path'] == [str(path_baboon_x4)]
np.testing.assert_almost_equal(results['ori_lq'][0], img_baboon_x4)
with pytest.raises(TypeError):
# filepath should be list
results = dict(lq_path=path_baboon_x4)
image_loader(results)
# convert to y-channel (bgr2y)
results = dict(lq_path=[str(path_baboon_x4), str(path_baboon)])
config = dict(io_backend='disk', key='lq', convert_to='y')
image_loader = LoadImageFromFileList(**config)
results = image_loader(results)
img_baboon_x4_y = mmcv.bgr2ycbcr(img_baboon_x4, y_only=True)
img_baboon_y = mmcv.bgr2ycbcr(img_baboon, y_only=True)
img_baboon_x4_y = np.expand_dims(img_baboon_x4_y, axis=2)
img_baboon_y = np.expand_dims(img_baboon_y, axis=2)
np.testing.assert_almost_equal(results['lq'][0], img_baboon_x4_y)
np.testing.assert_almost_equal(results['lq'][1], img_baboon_y)
assert results['lq_path'] == [str(path_baboon_x4), str(path_baboon)]
# convert to y-channel (rgb2y)
results = dict(lq_path=[str(path_baboon_x4), str(path_baboon)])
config = dict(io_backend='disk',
key='lq',
channel_order='rgb',
convert_to='y')
image_loader = LoadImageFromFileList(**config)
results = image_loader(results)
np.testing.assert_almost_equal(results['lq'][0], img_baboon_x4_y)
np.testing.assert_almost_equal(results['lq'][1], img_baboon_y)
assert results['lq_path'] == [str(path_baboon_x4), str(path_baboon)]
# convert to y-channel (ValueError)
results = dict(lq_path=[str(path_baboon_x4), str(path_baboon)])
config = dict(io_backend='disk', key='lq', convert_to='abc')
image_loader = LoadImageFromFileList(**config)
with pytest.raises(ValueError):
results = image_loader(results)
