【Pytorch】Pytorch工具轮子代码

Pytorch轮子子子子子子

加载图片转化为Tenosr格式

from torchvision import transforms
from PIL import Image
import torch

def image_loader(image_name):
  image = Image.open(image_name).convert('RGB')
  image = image.resize((256, 256))
  loader = transforms.Compose([transforms.ToTensor()]) 
  device = torch.device('cuda')
  image = loader(image).unsqueeze(0)/255.  #归一化处理
  return image.to(device, torch.float)


img1 = image_loader("einstein.png")
if torch.cuda.is_available():
    img1 = img1.cuda()

计算PSNR和SSIM

import cv2
import os
from skimage.measure import compare_ssim

sr_dir = os.listdir('../dataset/gt')
hr_dir = os.listdir('../dataset/haze')
list_len = len(sr_dir)
psnr = 0.0
ssim = 0.0
n = 0

def to_grey(image):
    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

for hr_image in hr_dir:
    for sr_image in sr_dir:
        if sr_image == hr_image:

            gt_img_path = os.path.join('../dataset/gt',sr_image)
            haze_img_path = os.path.join('../dataset/haze',hr_image)
 
            compute_psnr = cv2.PSNR(cv2.imread(gt_img_path), cv2.imread(haze_img_path))
            compute_ssim = compare_ssim(to_grey(cv2.imread(gt_img_path)),
                                        to_grey(cv2.imread(haze_img_path)))
            print(sr_image,"-psnr = ", compute_psnr)
            print(sr_image,"-ssim = ", compute_ssim)
            print("=====================")
            psnr += compute_psnr
            ssim += compute_ssim
            n += 1
            if n%list_len == 0:
                print("finish compute [%d/%d]" % (n, len(hr_dir)))

psnr = psnr / n
ssim = ssim / n
print("average psnr = ", psnr)
print("average ssim = ", ssim)