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)

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章