利用保存好的图片进行PSNR测试

from util_ import calc_psnr_and_ssim
import torch
import numpy as np
from glob import glob #import glob是错的
import os
from PIL import Image
from torchvision import transforms
# path = './data/test/MIDDLEBURY/'
# file_list = sorted(glob(os.path.join(path,'*L.png')))
# file_list_ref= sorted(glob(os.path.join(path,'*R.png')))
# for i in range(len(file_list)):
#     print(file_list[i])
#     print(file_list_ref[i])
#
trans = transforms.Compose([
    # transforms.Resize((256, 256)),
    transforms.ToTensor(),
    # transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
psnr_list = []

ssim_list = []
# sr_list = sorted(glob(
#             os.path.join('/home/cl/DRN_zxy/DRN-master/experiments/results/epoch800/%s/x4'%datasets[j],'*L.*')))
# hr_list = sorted(glob(
#             os.path.join('/home/cl/DRN_zxy/DRN-master/srdata/benchmark/%s/HR'%datasets[j],'*L.*')))
sr_list = sorted(glob(
            os.path.join('/root/userfolder/SRNTT-master/SRNTT-master/SR','*.png')))

hr_list = sorted(glob(
            os.path.join('/root/userfolder/SRNTT-master/SRNTT-master/data/test/Set13','*.png')))

for i in range(len(hr_list)):

    path_sr = sr_list[i]
    path_hr = hr_list[i]
    print(path_sr)
    print(path_hr)
    SR_left = Image.open(path_sr).convert('RGB')

    HR_left = Image.open(path_hr).convert('RGB')

    SR_left = trans(SR_left)
    HR_left = trans(HR_left)
    # print('SR',SR_left)
    # print('HR',HR_left)
    SR_left = torch.unsqueeze(SR_left,0)

    HR_left = torch.unsqueeze(HR_left, 0)
    psnr, ssim  = calc_psnr_and_ssim(SR_left.detach(), HR_left.detach())
    print(psnr,ssim)
    # SR_left = torch.unsqueeze(SR_left,0)



    # HR_left = torch.unsqueeze(HR_left, 0)

    psnr_list.append(psnr)

    ssim_list.append(ssim)
# print( 'mean psnr:',float(np.array(psnr_list).mean()))
# print( 'mean ssim:',float(np.array(ssim_list).mean()))
print( 'mean psnr:',float(np.array(psnr_list).mean()))
print( 'mean ssim:',float(np.array(ssim_list).mean()))
import math
import numpy as np
import logging
import cv2
import os
import shutil

import torch
import torch.nn as nn
import torch.nn.functional as F
def calc_psnr(img1, img2):
    ### args:
    # img1: [h, w, c], range [0, 255]
    # img2: [h, w, c], range [0, 255]
    diff = (img1 - img2) / 255.0
    diff[:, :, 0] = diff[:, :, 0] * 65.738 / 256.0
    diff[:, :, 1] = diff[:, :, 1] * 129.057 / 256.0
    diff[:, :, 2] = diff[:, :, 2] * 25.064 / 256.0

    diff = np.sum(diff, axis=2)
    mse = np.mean(np.power(diff, 2))
    return -10 * math.log10(mse)


def calc_ssim(img1, img2):
    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()

    ### args:
    # img1: [h, w, c], range [0, 255]
    # img2: [h, w, c], range [0, 255]
    # the same outputs as MATLAB's
    border = 0
    img1_y = np.dot(img1, [65.738, 129.057, 25.064]) / 256.0 + 16.0
    img2_y = np.dot(img2, [65.738, 129.057, 25.064]) / 256.0 + 16.0
    if not img1.shape == img2.shape:
        raise ValueError('Input images must have the same dimensions.')
    h, w = img1.shape[:2]
    img1_y = img1_y[border:h - border, border:w - border]
    img2_y = img2_y[border:h - border, border:w - border]

    if img1_y.ndim == 2:
        return ssim(img1_y, img2_y)
    elif img1.ndim == 3:
        if img1.shape[2] == 3:
            ssims = []
            for i in range(3):
                ssims.append(ssim(img1, img2))
            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 calc_psnr_and_ssim(sr, hr):
    ### args:
    # sr: pytorch tensor, range [-1, 1]
    # hr: pytorch tensor, range [-1, 1]

    ### prepare data
    # sr = (sr + 1.) * 127.5
    # hr = (hr + 1.) * 127.5
    sr = sr*255.
    hr = hr*255.
    if (sr.size() != hr.size()):
        h_min = min(sr.size(2), hr.size(2))
        w_min = min(sr.size(3), hr.size(3))
        sr = sr[:, :, :h_min, :w_min]
        hr = hr[:, :, :h_min, :w_min]

    img1 = np.transpose(sr.squeeze().round().cpu().numpy(), (1, 2, 0))
    img2 = np.transpose(hr.squeeze().round().cpu().numpy(), (1, 2, 0))

    psnr = calc_psnr(img1, img2)
    ssim = calc_ssim(img1, img2)

    return psnr, ssim

 

具体的计算方式可以自己选择,方法统一就好

 

posted on 2021-05-15 15:33  cltt  阅读(354)  评论(0编辑  收藏  举报

导航