NITIRE 2023官方的PSNR及SSIM计算代码

问题描述

做图像复原任务时，总避免不了计算PSNR和SSIM等图像质量评估指标，但是网上实在是太多计算这类指标的代码了，不同代码计算的结果还可能存在差异。有使用matlab计算SSIM的，也有使用python计算SSIM的。有使用第三方库skimage计算SSIM的，也有使用第三方库pyiqa计算SSIM的。博主曾经试过，对Urban100测试集进行SSIM的计算，即计算SR图像与HR图像的SSIM时，数值差异较大，使用matlab计算的SSIM为0.9644，而使用python计算的SSIM为0.7980。因此，这让人一头雾水，不知道哪个才是所谓的标准。

NITIRE比赛

博主关注了今年的NITIRE 2023超分赛道的比赛，留意到轻量级超分赛道NTIRE 2023 Challenge on Efficient Super-Resolution中，提供了一个代码模板，让参赛者可以自己测试模型的显存占用，推理时间，PSNR和SSIM。报告原文有句话是这么说的：

A code example for calculating these metrics is available at https://github.com/ofsoundof/NTIRE2023_ESR. The code of the submitted solutions and the pre-trained weights are also available in this repository.

计算这些指标的代码例子可以在https://github.com/ofsoundof/NTIRE2023_ESR中找到。所提交的方案的代码和预训练权重都能在这个仓库中找到。

也就是我们可以在https://github.com/ofsoundof/NTIRE2023_ESR中找到NITIRE官方计算PSNR和SSIM的代码，由于NITRE比赛是CVPR(IEEE Conference on Computer Vision and Pattern Recognition)举办的极具影响力的计算机视觉底层任务比赛，十分具有公信力，因此建议使用此代码进行计算。

核心代码

计算PSNR和SSIM的核心代码

'''
# =======================================
# metric, PSNR and SSIM
# =======================================
'''

# ----------
# PSNR
# ----------
def calculate_psnr(img1, img2, border=0):
    img1 = _bord_img(img1)
    img2 = _bord_img(img2)
    return _calculate_psnr(img1, img2)


# ----------
# SSIM
# ----------
def calculate_ssim(img1, img2, border=0):
    img1 = _bord_img(img1)
    img2 = _bord_img(img2)
    return _calculate_ssim(img1, img2)


def _calculate_ssim(img, img2, test_y_channel=True):
    if test_y_channel:
        img = to_y_channel(img)
        img2 = to_y_channel(img2)

    ssims = []
    for i in range(img.shape[2]):
        ssims.append(_ssim(img[..., i], img2[..., i]))
    return np.array(ssims).mean()

# 在y通道上计算
def _calculate_psnr(img, img2, test_y_channel=True,):
    if test_y_channel:
        img = to_y_channel(img)
        img2 = to_y_channel(img2)

    mse = np.mean((img - img2)**2)
    if mse == 0:
        return float('inf')
    return 20. * np.log10(255. / np.sqrt(mse))


def _ssim(img, img2):
    c1 = (0.01 * 255)**2
    c2 = (0.03 * 255)**2

    img = img.astype(np.float64)
    img2 = img2.astype(np.float64)
    kernel = cv2.getGaussianKernel(11, 1.5)
    window = np.outer(kernel, kernel.transpose())

    mu1 = cv2.filter2D(img, -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(img**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(img * 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()