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
具体的计算方式可以自己选择,方法统一就好