psnr ssim lpips

import os
import torch
import torchvision.transforms.functional as TF
from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.metrics import structural_similarity as ssim
import lpips
from PIL import Image

def calculate_metrics(folder1, folder2):
psnr_values = []
ssim_values = []
lpips_loss = lpips.LPIPS(net='vgg').cuda()
lpips_loss.eval()
lpips_values = []

for file1, file2 in zip(sorted(os.listdir(folder1)), sorted(os.listdir(folder2))):
image1 = Image.open(os.path.join(folder1, file1))
image2 = Image.open(os.path.join(folder2, file2))

# Convert images to tensors
image1 = TF.to_tensor(image1).unsqueeze(0).cuda()
image2 = TF.to_tensor(image2).unsqueeze(0).cuda()

# PSNR
psnr_val = psnr(image1.squeeze().cpu().numpy(), image2.squeeze().cpu().numpy())
psnr_values.append(psnr_val)

# SSIM
ssim_val = ssim(image1.squeeze().cpu().numpy(), image2.squeeze().cpu().numpy(), multichannel=True)
ssim_values.append(ssim_val)

# LPIPS
lpips_val = lpips_loss(image1, image2).item()
lpips_values.append(lpips_val)

avg_psnr = sum(psnr_values) / len(psnr_values)
avg_ssim = sum(ssim_values) / len(ssim_values)
avg_lpips = sum(lpips_values) / len(lpips_values)

return avg_psnr, avg_ssim, avg_lpips

folder1 = 'folder1_path'
folder2 = 'folder2_path'

psnr_val, ssim_val, lpips_val = calculate_metrics(folder1, folder2)
print(f'Average PSNR: {psnr_val}')
print(f'Average SSIM: {ssim_val}')
print(f'Average LPIPS: {lpips_val}')