NITIRE 2023 official PSNR and SSIM calculation code

NITIRE 2023 official PSNR and SSIM calculation code

Problem Description

When doing image restoration tasks, it is inevitable to calculate image quality evaluation indicators such as sum PSNRand SSIMsum. However, there are too many codes for calculating such indicators on the Internet, and the calculation results of different codes may be different. Some are calculated using matlabSSIM , and some are calculated using pythonSSIM . Some use third-party libraries skimageto calculate SSIM, and some use third-party libraries pyiqato calculate SSIM. The blogger once tried to calculate Urban100the test set SSIM, that is, when calculating the SR image and the HR image SSIM, the numerical difference is large. The calculation using matlabSSIM is 0.9644, while the calculation using pythonSSIM is 0.7980. So it leaves people confused as to which one is the so-called standard .

NITIRE competition

The blogger paid attention to this year's NITIRE 2023 super-resolution track competition and noticed that the lightweight super-resolution track NTIRE 2023 Challenge on Efficient Super-Resolution provides a code template so that contestants can test the video memory usage of the model by themselves . , inference time , PSNR and SSIM . There is a sentence in the original report that says:

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.

Code examples for calculating these metrics can be found at https://github.com/ofsoundof/NTIRE2023_ESR . The code and pre-trained weights of the submitted solution can be found in this repository.

That is to say, we can find the NITIRE official calculation code at https://github.com/ofsoundof/NTIRE2023_ESR . Since the NITRE competition is a very influential computer vision bottom layer organized by CVPR (IEEE Conference on Computer Vision and Pattern Recognition) The task competition is very credible, so it is recommended to use this code for calculations.PSNRSSIM

core code

Core code for calculating PSNR and 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()

Finally, thank you friends for your study~

Finally, a link to the 2023 Efficient Super-Resolution competition report is attached. Everyone is welcome to read and share: NTIRE 2023 Challenge on Efficient Super-Resolution: Methods and Results