某次圖像生成實驗的保存結果是下面這些圖片,現在想要計算相同迭代次數下的*_real_image*和*_synthesized*兩張圖片之間的SSIM和PSNR。
整理出的代碼如下:
import numpy
import numpy as np
import math
import cv2
import torch
import pytorch_ssim
from torch.autograd import Variable
import os
import re
def psnr(img1, img2):
mse = numpy.mean( (img1 - img2) ** 2 )
if mse == 0:
return 100
PIXEL_MAX = 255.0
return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
def ssim(img1,img2):
img1 = torch.from_numpy(np.rollaxis(img1, 2)).float().unsqueeze(0)/255.0
img2 = torch.from_numpy(np.rollaxis(img2, 2)).float().unsqueeze(0)/255.0
img1 = Variable( img1, requires_grad=False) # torch.Size([256, 256, 3])
img2 = Variable( img2, requires_grad = False)
ssim_value = pytorch_ssim.ssim(img1, img2).item()
return ssim_value
dir = "./images/"
for item in os.listdir(dir):
# 首先判斷圖片類別
# 如果是real圖片,則找到與其對應的生成圖片,並進行計算;否則,判斷下一個文件。
if re.search('real',item) != None: # 是real圖片
realImagePath = os.path.join(dir,item) # './images/epoch080_iter1590000_real_image_3983_1.png'
try: # 爲了防止意外找不到文件的情況
realImage = cv2.imread(realImagePath) # yes
# 根據這張圖片的文件名確定它對應的生成圖片的文件名
# 使用正則表達式中的sub替換函數
itemSyn = re.sub("real","synthesized",item)
synImagePath = os.path.join(dir,itemSyn)
synImage = cv2.imread(synImagePath)
psnrValue = psnr(realImage,synImage)
ssimValue = ssim(realImage,synImage)
print('psnr:' + str(psnrValue))
print('ssim' + str(ssimValue))
except: # 如果無該文件
print('ERROR! can not find' + item)
else:
print('success:' + item)