2020年3月13日星期五

Python regular expressions match specific input and output images

The saved results of an image generation experiment are the following pictures. Now I want to calculate the SSIM and PSNR between * _real_image * and * _synthesized * in the same iterations.






The sorted out code is as follows:

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):
    # First determine the picture category
    # If it is a real picture, the corresponding generated picture is found and calculated; otherwise, the next file is judged.
    if re.search('real',item) != None:  # 是real图片
        realImagePath = os.path.join(dir,item)  # './images/epoch080_iter1590000_real_image_3983_1.png'
        try:    # To prevent accidentally missing files
            realImage = cv2.imread(realImagePath)   # yes
            # Determine the file name of the corresponding image based on the file name of this image
            # Replace function with sub in regular expression
            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: # Without the file
            print('ERROR! can not find' + item)
        else:
            print('success:' + item)
 

没有评论:

发表评论