計算lpips

helloWorldhelloWorld發表於2024-03-14
import os
import torch
import lpips
from PIL import Image
from torchvision.transforms import ToTensor

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

def calculate_lpips(img_path1, img_path2):
    # 讀取影像
    img1 = Image.open(img_path1)
    img2 = Image.open(img_path2)

    # 將影像轉換為張量
    img1_tensor = ToTensor()(img1).unsqueeze(0)
    img2_tensor = ToTensor()(img2).unsqueeze(0)

    # 計算LPIPS
    lpips_value = loss_fn(img1_tensor, img2_tensor)

    return lpips_value.item()

# 資料夾路徑
folder_path1 = 'your_folder_path1'
folder_path2 = 'your_folder_path2'

# 獲取資料夾中的所有影像路徑
image_paths1 = [os.path.join(folder_path1, img) for img in os.listdir(folder_path1) if img.endswith('.jpg')]
image_paths2 = [os.path.join(folder_path2, img) for img in os.listdir(folder_path2) if img.endswith('.jpg')]

# 確保兩個資料夾中有相同名稱的影像
assert set(os.path.basename(p) for p in image_paths1) == set(os.path.basename(p) for p in image_paths2), "The two folders must contain images with the same names."

# 計算並列印所有同名影像對的LPIPS
lpips_values = []
for img_name in os.listdir(folder_path1):
    if img_name.endswith('.jpg'):
        img_path1 = os.path.join(folder_path1, img_name)
        img_path2 = os.path.join(folder_path2, img_name)
        lpips_value = calculate_lpips(img_path1, img_path2)
        lpips_values.append(lpips_value)
        print(f'LPIPS between {img_path1} and {img_path2}: {lpips_value}')

# 計算並列印LPIPS的平均值
average_lpips = sum(lpips_values) / len(lpips_values)
print(f'Average LPIPS: {average_lpips}')

相關文章