Spaces:
Runtime error
Runtime error
| # ----------------------------------------------------------------------------------- | |
| # https://github.com/JingyunLiang/SwinIR/blob/main/utils/util_calculate_psnr_ssim.py | |
| # ----------------------------------------------------------------------------------- | |
| 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) | |
| # follow https://gitlab.com/Queuecumber/quantization-guided-ac/-/blob/master/metrics/psnrb.py | |
| 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 |