Python正则表达式匹配特定的的输入输出图片

     某次图像生成实验的保存结果是下面这些图片,现在想要计算相同迭代次数下的*_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)

发布了68 篇原创文章 · 获赞 12 · 访问量 2万+

猜你喜欢

转载自blog.csdn.net/qazwsxrx/article/details/104559781