|
|
|
|
|
|
|
|
|
|
|
|
|
|
import cv2
|
|
|
import torch
|
|
|
import numpy as np
|
|
|
|
|
|
def calculate_psnr(img1, img2, crop_border, input_order='HWC', test_y_channel=False):
|
|
|
"""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 edge of an image. These
|
|
|
pixels are not involved in the PSNR calculation.
|
|
|
input_order (str): Whether the input order is 'HWC' or 'CHW'.
|
|
|
Default: 'HWC'.
|
|
|
test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
|
|
|
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)
|
|
|
img1 = img1.astype(np.float64)
|
|
|
img2 = img2.astype(np.float64)
|
|
|
|
|
|
if crop_border != 0:
|
|
|
img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
|
|
|
img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
|
|
|
|
|
|
if test_y_channel:
|
|
|
img1 = to_y_channel(img1)
|
|
|
img2 = to_y_channel(img2)
|
|
|
|
|
|
mse = np.mean((img1 - img2) ** 2)
|
|
|
if mse == 0:
|
|
|
return float('inf')
|
|
|
return 20. * np.log10(255. / np.sqrt(mse))
|
|
|
|
|
|
|
|
|
def _ssim(img1, img2):
|
|
|
"""Calculate SSIM (structural similarity) for one channel images.
|
|
|
It is called by func:`calculate_ssim`.
|
|
|
Args:
|
|
|
img1 (ndarray): Images with range [0, 255] with order 'HWC'.
|
|
|
img2 (ndarray): Images with range [0, 255] with order 'HWC'.
|
|
|
Returns:
|
|
|
float: ssim result.
|
|
|
"""
|
|
|
|
|
|
C1 = (0.01 * 255) ** 2
|
|
|
C2 = (0.03 * 255) ** 2
|
|
|
|
|
|
img1 = img1.astype(np.float64)
|
|
|
img2 = img2.astype(np.float64)
|
|
|
kernel = cv2.getGaussianKernel(11, 1.5)
|
|
|
window = np.outer(kernel, kernel.transpose())
|
|
|
|
|
|
mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
|
|
|
mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
|
|
|
mu1_sq = mu1 ** 2
|
|
|
mu2_sq = mu2 ** 2
|
|
|
mu1_mu2 = mu1 * mu2
|
|
|
sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
|
|
|
sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
|
|
|
sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
|
|
|
|
|
|
ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
|
|
|
return ssim_map.mean()
|
|
|
|
|
|
|
|
|
def calculate_ssim(img1, img2, crop_border, input_order='HWC', test_y_channel=False):
|
|
|
"""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 edge of an image. These
|
|
|
pixels are not involved in the SSIM calculation.
|
|
|
input_order (str): Whether the input order is 'HWC' or 'CHW'.
|
|
|
Default: 'HWC'.
|
|
|
test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
|
|
|
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)
|
|
|
img1 = img1.astype(np.float64)
|
|
|
img2 = img2.astype(np.float64)
|
|
|
|
|
|
if crop_border != 0:
|
|
|
img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
|
|
|
img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
|
|
|
|
|
|
if test_y_channel:
|
|
|
img1 = to_y_channel(img1)
|
|
|
img2 = to_y_channel(img2)
|
|
|
|
|
|
ssims = []
|
|
|
for i in range(img1.shape[2]):
|
|
|
ssims.append(_ssim(img1[..., i], img2[..., i]))
|
|
|
return np.array(ssims).mean()
|
|
|
|
|
|
|
|
|
def _blocking_effect_factor(im):
|
|
|
block_size = 8
|
|
|
|
|
|
block_horizontal_positions = torch.arange(7, im.shape[3] - 1, 8)
|
|
|
block_vertical_positions = torch.arange(7, im.shape[2] - 1, 8)
|
|
|
|
|
|
horizontal_block_difference = (
|
|
|
(im[:, :, :, block_horizontal_positions] - im[:, :, :, block_horizontal_positions + 1]) ** 2).sum(
|
|
|
3).sum(2).sum(1)
|
|
|
vertical_block_difference = (
|
|
|
(im[:, :, block_vertical_positions, :] - im[:, :, block_vertical_positions + 1, :]) ** 2).sum(3).sum(
|
|
|
2).sum(1)
|
|
|
|
|
|
nonblock_horizontal_positions = np.setdiff1d(torch.arange(0, im.shape[3] - 1), block_horizontal_positions)
|
|
|
nonblock_vertical_positions = np.setdiff1d(torch.arange(0, im.shape[2] - 1), block_vertical_positions)
|
|
|
|
|
|
horizontal_nonblock_difference = (
|
|
|
(im[:, :, :, nonblock_horizontal_positions] - im[:, :, :, nonblock_horizontal_positions + 1]) ** 2).sum(
|
|
|
3).sum(2).sum(1)
|
|
|
vertical_nonblock_difference = (
|
|
|
(im[:, :, nonblock_vertical_positions, :] - im[:, :, nonblock_vertical_positions + 1, :]) ** 2).sum(
|
|
|
3).sum(2).sum(1)
|
|
|
|
|
|
n_boundary_horiz = im.shape[2] * (im.shape[3] // block_size - 1)
|
|
|
n_boundary_vert = im.shape[3] * (im.shape[2] // block_size - 1)
|
|
|
boundary_difference = (horizontal_block_difference + vertical_block_difference) / (
|
|
|
n_boundary_horiz + n_boundary_vert)
|
|
|
|
|
|
n_nonboundary_horiz = im.shape[2] * (im.shape[3] - 1) - n_boundary_horiz
|
|
|
n_nonboundary_vert = im.shape[3] * (im.shape[2] - 1) - n_boundary_vert
|
|
|
nonboundary_difference = (horizontal_nonblock_difference + vertical_nonblock_difference) / (
|
|
|
n_nonboundary_horiz + n_nonboundary_vert)
|
|
|
|
|
|
scaler = np.log2(block_size) / np.log2(min([im.shape[2], im.shape[3]]))
|
|
|
bef = scaler * (boundary_difference - nonboundary_difference)
|
|
|
|
|
|
bef[boundary_difference <= nonboundary_difference] = 0
|
|
|
return bef
|
|
|
|
|
|
|
|
|
def calculate_psnrb(img1, img2, crop_border, input_order='HWC', test_y_channel=False):
|
|
|
"""Calculate PSNR-B (Peak Signal-to-Noise Ratio).
|
|
|
Ref: Quality assessment of deblocked images, for JPEG image deblocking evaluation
|
|
|
# https://gitlab.com/Queuecumber/quantization-guided-ac/-/blob/master/metrics/psnrb.py
|
|
|
Args:
|
|
|
img1 (ndarray): Images with range [0, 255].
|
|
|
img2 (ndarray): Images with range [0, 255].
|
|
|
crop_border (int): Cropped pixels in each edge of an image. These
|
|
|
pixels are not involved in the PSNR calculation.
|
|
|
input_order (str): Whether the input order is 'HWC' or 'CHW'.
|
|
|
Default: 'HWC'.
|
|
|
test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
|
|
|
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)
|
|
|
img1 = img1.astype(np.float64)
|
|
|
img2 = img2.astype(np.float64)
|
|
|
|
|
|
if crop_border != 0:
|
|
|
img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
|
|
|
img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
|
|
|
|
|
|
if test_y_channel:
|
|
|
img1 = to_y_channel(img1)
|
|
|
img2 = to_y_channel(img2)
|
|
|
|
|
|
|
|
|
img1 = torch.from_numpy(img1).permute(2, 0, 1).unsqueeze(0) / 255.
|
|
|
img2 = torch.from_numpy(img2).permute(2, 0, 1).unsqueeze(0) / 255.
|
|
|
|
|
|
total = 0
|
|
|
for c in range(img1.shape[1]):
|
|
|
mse = torch.nn.functional.mse_loss(img1[:, c:c + 1, :, :], img2[:, c:c + 1, :, :], reduction='none')
|
|
|
bef = _blocking_effect_factor(img1[:, c:c + 1, :, :])
|
|
|
|
|
|
mse = mse.view(mse.shape[0], -1).mean(1)
|
|
|
total += 10 * torch.log10(1 / (mse + bef))
|
|
|
|
|
|
return float(total) / img1.shape[1]
|
|
|
|
|
|
|
|
|
def reorder_image(img, input_order='HWC'):
|
|
|
"""Reorder images to 'HWC' order.
|
|
|
If the input_order is (h, w), return (h, w, 1);
|
|
|
If the input_order is (c, h, w), return (h, w, c);
|
|
|
If the input_order is (h, w, c), return as it is.
|
|
|
Args:
|
|
|
img (ndarray): Input image.
|
|
|
input_order (str): Whether the input order is 'HWC' or 'CHW'.
|
|
|
If the input image shape is (h, w), input_order will not have
|
|
|
effects. Default: 'HWC'.
|
|
|
Returns:
|
|
|
ndarray: reordered image.
|
|
|
"""
|
|
|
|
|
|
if input_order not in ['HWC', 'CHW']:
|
|
|
raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' "'HWC' and 'CHW'")
|
|
|
if len(img.shape) == 2:
|
|
|
img = img[..., None]
|
|
|
if input_order == 'CHW':
|
|
|
img = img.transpose(1, 2, 0)
|
|
|
return img
|
|
|
|
|
|
|
|
|
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 = bgr2ycbcr(img, y_only=True)
|
|
|
img = img[..., None]
|
|
|
return img * 255.
|
|
|
|
|
|
|
|
|
def _convert_input_type_range(img):
|
|
|
"""Convert the type and range of the input image.
|
|
|
It converts the input image to np.float32 type and range of [0, 1].
|
|
|
It is mainly used for pre-processing the input image in colorspace
|
|
|
convertion functions such as rgb2ycbcr and ycbcr2rgb.
|
|
|
Args:
|
|
|
img (ndarray): The input image. It accepts:
|
|
|
1. np.uint8 type with range [0, 255];
|
|
|
2. np.float32 type with range [0, 1].
|
|
|
Returns:
|
|
|
(ndarray): The converted image with type of np.float32 and range of
|
|
|
[0, 1].
|
|
|
"""
|
|
|
img_type = img.dtype
|
|
|
img = img.astype(np.float32)
|
|
|
if img_type == np.float32:
|
|
|
pass
|
|
|
elif img_type == np.uint8:
|
|
|
img /= 255.
|
|
|
else:
|
|
|
raise TypeError('The img type should be np.float32 or np.uint8, ' f'but got {img_type}')
|
|
|
return img
|
|
|
|
|
|
|
|
|
def _convert_output_type_range(img, dst_type):
|
|
|
"""Convert the type and range of the image according to dst_type.
|
|
|
It converts the image to desired type and range. If `dst_type` is np.uint8,
|
|
|
images will be converted to np.uint8 type with range [0, 255]. If
|
|
|
`dst_type` is np.float32, it converts the image to np.float32 type with
|
|
|
range [0, 1].
|
|
|
It is mainly used for post-processing images in colorspace convertion
|
|
|
functions such as rgb2ycbcr and ycbcr2rgb.
|
|
|
Args:
|
|
|
img (ndarray): The image to be converted with np.float32 type and
|
|
|
range [0, 255].
|
|
|
dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
|
|
|
converts the image to np.uint8 type with range [0, 255]. If
|
|
|
dst_type is np.float32, it converts the image to np.float32 type
|
|
|
with range [0, 1].
|
|
|
Returns:
|
|
|
(ndarray): The converted image with desired type and range.
|
|
|
"""
|
|
|
if dst_type not in (np.uint8, np.float32):
|
|
|
raise TypeError('The dst_type should be np.float32 or np.uint8, ' f'but got {dst_type}')
|
|
|
if dst_type == np.uint8:
|
|
|
img = img.round()
|
|
|
else:
|
|
|
img /= 255.
|
|
|
return img.astype(dst_type)
|
|
|
|
|
|
|
|
|
def bgr2ycbcr(img, y_only=False):
|
|
|
"""Convert a BGR image to YCbCr image.
|
|
|
The bgr version of rgb2ycbcr.
|
|
|
It implements the ITU-R BT.601 conversion for standard-definition
|
|
|
television. See more details in
|
|
|
https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
|
|
|
It differs from a similar function in cv2.cvtColor: `BGR <-> YCrCb`.
|
|
|
In OpenCV, it implements a JPEG conversion. See more details in
|
|
|
https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
|
|
|
Args:
|
|
|
img (ndarray): The input image. It accepts:
|
|
|
1. np.uint8 type with range [0, 255];
|
|
|
2. np.float32 type with range [0, 1].
|
|
|
y_only (bool): Whether to only return Y channel. Default: False.
|
|
|
Returns:
|
|
|
ndarray: The converted YCbCr image. The output image has the same type
|
|
|
and range as input image.
|
|
|
"""
|
|
|
img_type = img.dtype
|
|
|
img = _convert_input_type_range(img)
|
|
|
if y_only:
|
|
|
out_img = np.dot(img, [24.966, 128.553, 65.481]) + 16.0
|
|
|
else:
|
|
|
out_img = np.matmul(
|
|
|
img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786], [65.481, -37.797, 112.0]]) + [16, 128, 128]
|
|
|
out_img = _convert_output_type_range(out_img, img_type)
|
|
|
return out_img |
