详细内容可参考: https://www.computationalimaging.cn/2020/03/python-regular-expressions-match.html
某次图像生成实验的保存结果是下面这些图片,现在想要计算相同迭代次数下的*_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)