评价指标SSIM与PSNR

之前也有对 超分前后的图像进行 评估结果的博客，但当时的psnr以及ssim函数都是自己查了一些资料写的，ssim的函数可能出了一些问题，导致结果偏低。
本次的psnr以及ssim都时调用ski-image库的，直接使用即可。

文件的目录结构如图所示：

import os
import numpy as np
from PIL import Image
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import peak_signal_noise_ratio as psnr


def evalution(generate_dir, gt_dir):
    # hr result's file list,for example:[001,003,005,006,007,008,013,017]
    generate_file_list = os.listdir(generate_dir)
    gt_file_list = os.listdir(gt_dir)

    fc = open(os.path.join(generate_dir, 'result.txt'), 'w')

    ssim_file_list = []
    psnr_file_list = []
    for file in generate_file_list:
        # cat file's path and file's name can get full path of file
        generate_file_path = os.path.join(generate_dir, file)
        gt_file_path = os.path.join(gt_dir, file)
        # print(generate_file_path) #E:\VSR\实验推理结果数据\ESPCN\test_amy\001

        # pic  list of file,for example:[00000000.png,00000001.png,……,00000029.png]
        generate_pic_list = os.listdir(generate_file_path)
        gt_pic_list = os.listdir(gt_file_path)

        ssim_pic_list = []
        psnr_pic_list = []
        for pic in generate_pic_list:
            # cat pic's path and pic's name can get full path of pic
            generate_pic_path = os.path.join(generate_file_path, pic)
            pic = pic.split('_')[1] #在tof时使用
            # print(pic)
            gt_pic_path = os.path.join(gt_file_path, pic)
            img1 = np.array(Image.open(generate_pic_path))
            img2 = np.array(Image.open(gt_pic_path))

            assert img1.shape == img2.shape, (
                f'Image shapes are different: {img1.shape}, {img2.shape}.')

            s = ssim(img1, img2, multichannel=True)
            ssim_pic_list.append(s)

            p = psnr(img1, img2)
            psnr_pic_list.append(p)
            # print('        ', pic, '---ssim:', s, '  psnr:', p)
            fc.write('        %s---ssim:%.4f  psnr:%.4f \n' % (pic, p, s))
        avg_ssim = np.mean(ssim_pic_list)
        avg_psnr = np.mean(psnr_pic_list)
        # print('   ', file, '----avg_ssim:', avg_ssim, '  avg_psnr:', avg_psnr)
        fc.write('   %s---avg_ssim:%.4f  avg_psnr:%.4f \n' % (file, avg_ssim, avg_psnr))

        ssim_file_list.append(avg_ssim)
        psnr_file_list.append(avg_psnr)
    print(ssim_file_list)
    print(psnr_file_list)
    # fc.write(str(ssim_file_list))
    # fc.write('\n')
    # fc.write(str(psnr_file_list))
    ssim_total_avg = np.mean(ssim_file_list)
    psnr_total_avg = np.mean(psnr_file_list)
    print('Total_avg_ssim:%f  Total_avg_psnr:%f\n' % (ssim_total_avg, psnr_total_avg))
    fc.write('\n Total_avg_ssim:%f  Total_avg_psnr:%f \n' % (ssim_total_avg, psnr_total_avg))
    fc.close()


if __name__ == "__main__":
    gt_dir = r'E:\VSR\实验推理结果数据\三个消融实验结果\test_amy\GT'
    generate_dir = r'E:\VSR\实验推理结果数据\tof\tof_x4_reds_trans_4\results_tof4-12w'
    evalution(generate_dir, gt_dir)