计算两个文件夹中图片的PSNR和SSIM

注释:两个文件夹中图片的文件名要一致,格式、后缀都要一直,如果其他格式可以将.jpg改成.bmp
写来只是自己看,自己能用,所以跑不通您也别怪我菜


```bash

```bash

```python
import numpy as np
import os
import cv2
import math

from skimage.measure import compare_ssim, compare_psnr
from functools import partial

IMG_EXTENSIONS = [
    '.jpg', '.JPG', '.jpeg', '.JPEG',
    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]
def isImagefile(input):
    return any(input.endswith(ext) for ext in IMG_EXTENSIONS)


def quality_assess_SM3(inputdir, gtdir,resultdir):
    inputdir_list = os.listdir(inputdir)
    gtdir_list = os.listdir(gtdir)
    inputdir_list.sort()
    gtdir_list.sort()

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

    ssim_list = []
    psnr_list = []
    n_images = 0
    avg_p = 0
    avg_s = 0
    for name in gtdir_list:
        if name.endswith('.png'):
            n_images += 1
    print('number of images :',n_images)
    for name in gtdir_list:

        path1 = os.path.join(inputdir,name)
        path2 = os.path.join(gtdir,name)
        img1 = cv2.imread(path1)
        img2 = cv2.imread(path2)

        inputImg_gry = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
        GT_gry = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

        print(name+':')
        p = compare_psnr(img1, img2)
        avg_p += p/n_images
        print('     	p', p)
        psnr_list.append(p)

        s = compare_ssim(inputImg_gry,GT_gry)
        avg_s += s/n_images
        print('     	s', s)
        ssim_list.append(s)

        fc.write('%s : %.4f %.4f \n' %(name, p, s)) 
    print('Avg. psnr: %f ssim: %f\n' %(avg_p, avg_s))
    fc.write('Avg. psnr: %f ssim: %f\n' %(avg_p, avg_s))
    fc.close()

def dropext(fn): return os.path.splitext(fn)[0]

gtdir = '/home/shaodongliang/code/IBCLN_ca/datasets/Test_IBCLN_ceil/GT/'
inputdir = '/home/shaodongliang/code/IBCLN_ca/results/IBCLN/test_115/images/'
resultdir = './results'
quality_assess_SM3(inputdir, gtdir,resultdir)

print("____________________take notes:network_epoch_datasets_______________________________")


你可能感兴趣的:(python,深度学习)