Python计算Y通道或者RGB通道的PSNR_SSIM

一、 PSNR与SSIM简介

• 峰值信噪比(Peak Signal-to-noise Ratio, PSNR) 和结构相似性 (Structural Similarity, SSIM) 常作为图像超分辨重建任务主要的评价指标。
• PSNR度量图像压缩或信号重建图像与原始图像的像素差异，单位为dB，其值越大，则表明图像的失真度越小，计算公式如下：
  
  公式中，MAXi 为图像最大像素值（RGB图像，像点颜色数值为255）；MSE表示两幅尺寸为m×n图像I与图像K的均方误差值。
• SSIM是衡量两幅图像在噪声干扰、畸变或者失真时亮度(luminance)、对比度(constrast)以及结构的指标(structure)，取值范围为[0,1]，数值越接近1，失真程度越小。其公式如下：
  

二、代码实现

• 主要的代码源于 ：

• 添加：PSNR/SSIM计算结果，保存文件（txt），便于后期导入excel列表分析。

• 代码介绍：PSNR_SSIM效果和matlab计算的结果一致（本人没有验证）。

• 代码功能，能够计算图像的Y_channel 或者RGB_channel状态下的PSNR_SSIM结果

• 用法：将真实图片，生成图片分别放入两个文件夹，代码中有选择两种方式（only_Y, RGB），切换注释相应的代码行即可。

• 注意：代码中涉及相应的代码库，如果没有，记得安装。

'''
calculate the PSNR and SSIM.
same as MATLAB's results
'''
import os
import math
import numpy as np
import cv2
import glob
import os

def main():
    # Configurations

    # GT - Ground-truth;
    # Gen: Generated / Restored / Recovered images
    folder_GT = '/content/drive/MyDrive/Experiment/codes/26_PSNR_SSIM/All/groud_truth_png'
    folder_Gen = '/content/drive/MyDrive/Experiment/codes/26_PSNR_SSIM/All/1_ESPCN_png'
   
    #save the psnr and ssim score by txt
    PS_path = '/content/drive/MyDrive/Experiment/codes/26_PSNR_SSIM/Score'
    if not os.path.exists(PS_path):
        print('NO_path to sava the NIQE_Score,Making....')
        os.makedirs(PS_path)
    else:
        print("The NIQE_Score_path has existed")
    PS_txt = open(PS_path +'/score1.txt', 'a' )

    crop_border = 4  # same with scale
    suffix = ''  # suffix for Gen images
    test_Y = False  # True: test Y channel only; False: test RGB channels

    PSNR_all = []
    SSIM_all = []
    img_list = sorted(glob.glob(folder_GT + '/*'))

    if test_Y:
        print('Testing Y channel.')
    else:
        print('Testing RGB channels.')

    for i, img_path in enumerate(img_list):
        base_name = os.path.splitext(os.path.basename(img_path))[0]
        print(base_name)
        im_GT = cv2.imread(img_path) / 255.
        #不同格式图像
        # im_Gen = cv2.imread(os.path.join(folder_Gen, base_name + suffix + '.tif')) / 255.
        im_Gen = cv2.imread(os.path.join(folder_Gen, base_name + suffix + '.png')) / 255.

        if test_Y and im_GT.shape[2] == 3:  # evaluate on Y channel in YCbCr color space
            im_GT_in = bgr2ycbcr(im_GT)
            im_Gen_in = bgr2ycbcr(im_Gen)
        else:
            im_GT_in = im_GT
            im_Gen_in = im_Gen

        # crop borders
        if crop_border == 0:
            cropped_GT = im_GT_in
            cropped_Gen = im_Gen_in
        else:
            if im_GT_in.ndim == 3:
                cropped_GT = im_GT_in[crop_border:-crop_border, crop_border:-crop_border, :]
                cropped_Gen = im_Gen_in[crop_border:-crop_border, crop_border:-crop_border, :]
            elif im_GT_in.ndim == 2:
                cropped_GT = im_GT_in[crop_border:-crop_border, crop_border:-crop_border]
                cropped_Gen = im_Gen_in[crop_border:-crop_border, crop_border:-crop_border]
            else:
                raise ValueError('Wrong image dimension: {}. Should be 2 or 3.'.format(im_GT_in.ndim))

        #不同通道数（Y通道和RGB三个通道），需要更改
        # calculate PSNR and SSIM
        # PSNR = calculate_psnr(cropped_GT * 255, cropped_Gen * 255)
        PSNR = calculate_rgb_psnr(cropped_GT * 255, cropped_Gen * 255)

        SSIM = calculate_ssim(cropped_GT * 255, cropped_Gen * 255)
        print('{:3d} - {:25}. \tPSNR: {:.4f} dB, \tSSIM: {:.4f}'.format(
            i + 1, base_name, PSNR, SSIM))
        PSNR_all.append(PSNR)
        SSIM_all.append(SSIM)

        single_info = '[{}],PSNR(dB),{:.4f}, SSIM,{:.4f}'
        PS_txt.write(single_info.format(base_name,PSNR,SSIM))
        PS_txt.write("\n")

    Mean_format = 'Mean_PSNR: {:.4f}, Mean_SSIM: {:.4f}'
    Mean_PSNR =  sum(PSNR_all) / len(PSNR_all)
    Mean_SSIM =  sum(SSIM_all) / len(SSIM_all)
    print(Mean_format.format(Mean_PSNR,Mean_SSIM))
    
    PS_txt.write(Mean_format.format(Mean_PSNR,Mean_SSIM))
    PS_txt.write('\n')

def calculate_psnr(img1, img2):
    # img1 and img2 have range [0, 255]
    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    mse = np.mean((img1 - img2)**2)
    if mse == 0:
        return float('inf')
    return 20 * math.log10(255.0 / math.sqrt(mse))

def calculate_rgb_psnr(img1, img2):
    """calculate psnr among rgb channel, img1 and img2 have range [0, 255]
    """
    n_channels = np.ndim(img1)
    sum_psnr = 0
    for i in range(n_channels):
        this_psnr = calculate_psnr(img1[:,:,i], img2[:,:,i])
        sum_psnr += this_psnr
    return sum_psnr/n_channels

def ssim(img1, img2):
    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]  # valid
    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):
    '''calculate SSIM
    the same outputs as MATLAB's
    img1, img2: [0, 255]
    '''
    if not img1.shape == img2.shape:
        raise ValueError('Input images must have the same dimensions.')
    if img1.ndim == 2:
        return ssim(img1, img2)
    elif img1.ndim == 3:
        if img1.shape[2] == 3:
            ssims = []
            for i in range(img1.shape[2]):
                ssims.append(ssim(img1[..., i], img2[..., i]))
            return np.array(ssims).mean()
        elif img1.shape[2] == 1:
            return ssim(np.squeeze(img1), np.squeeze(img2))
    else:
        raise ValueError('Wrong input image dimensions.')

def bgr2ycbcr(img, only_y=True):
    '''same as matlab rgb2ycbcr
    only_y: only return Y channel
    Input:
        uint8, [0, 255]
        float, [0, 1]
    '''
    in_img_type = img.dtype
    img.astype(np.float32)
    if in_img_type != np.uint8:
        img *= 255.
    # convert
    if only_y:
        rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
    else:
        rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
                              [65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
    if in_img_type == np.uint8:
        rlt = rlt.round()
    else:
        rlt /= 255.
    return rlt.astype(in_img_type)


if __name__ == '__main__':
    main()

三、运行、保存结果展示

• 包括单张和平均的PSNR_SSIM结果