44

import os
import torch
import lpips
import numpy as np
from PIL import Image
from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.metrics import structural_similarity as ssim

# 初始化LPIPS模型
loss_fn = lpips.LPIPS(net='vgg')

def calculate_metrics(folder1, folder2, ext='.jpg'):
    # 获取文件夹中的文件名
    filenames = [f for f in os.listdir(folder1) if f.endswith(ext)]
    
    psnr_values = []
    ssim_values = []
    lpips_values = []
    
    for filename in filenames:
        # 加载图像
        img1 = Image.open(os.path.join(folder1, filename))
        img2 = Image.open(os.path.join(folder2, filename))
        
        # 转换为numpy数组
        img1_np = np.array(img1)
        img2_np = np.array(img2)
        
        # 计算PSNR
        psnr_value = psnr(img1_np, img2_np)
        
        # 计算SSIM
        ssim_value = ssim(img1_np, img2_np, multichannel=True, win_size=min(img1_np.shape[:2]))  
                                                                     # 使用图像的较小边作为窗口大小
        # 转换为张量
        img1_tensor = torch.from_numpy(img1_np).unsqueeze(0).float()
        img2_tensor = torch.from_numpy(img2_np).unsqueeze(0).float()
        
        # 计算LPIPS
        lpips_value = loss_fn(img1_tensor, img2_tensor)
        
        psnr_values.append(psnr_value)
        ssim_values.append(ssim_value)
        lpips_values.append(lpips_value.item())
    
    return psnr_values, ssim_values, lpips_values

# 使用函数
folder1 = 'path_to_folder1'
folder2 = 'path_to_folder2'
psnr_values, ssim_values, lpips_values = calculate_metrics(folder1, folder2)