1

import os
import cv2
import torch
import pytorch_msssim

# 定义函数用于计算两张图片的 SSIM 值
def calc_ssim(img1, img2):
    # 转换为 PyTorch Tensor, 并转为 float 类型
    img1_tensor = torch.from_numpy(img1).float().unsqueeze(0).permute(0, 3, 1, 2) / 255.0
    img2_tensor = torch.from_numpy(img2).float().unsqueeze(0).permute(0, 3, 1, 2) / 255.0

    # 计算 SSIM
    ssim_val = pytorch_msssim.ssim(img1_tensor, img2_tensor, data_range=1.0, size_average=True)

    return ssim_val.item()

# 定义函数用于计算文件夹下所有图片的 SSIM 平均值
def calc_ssim_folder(folder1, folder2):
    # 获取文件夹下所有图片的文件名
    filenames1 = os.listdir(folder1)
    filenames2 = os.listdir(folder2)

    # 检查两个文件夹中图片数量是否相同
    if len(filenames1) != len(filenames2):
        print("Error: 两个文件夹中图片数量不相同")
        return None

    # 计算所有图片的 SSIM 值
    ssim_sum = 0
    for filename1, filename2 in zip(filenames1, filenames2):
        # 读入图片
        img1 = cv2.imread(os.path.join(folder1, filename1))
        img2 = cv2.imread(os.path.join(folder2, filename2))

        # 计算 SSIM 值
        ssim_val = calc_ssim(img1, img2)

        # 累加 SSIM 值
        ssim_sum += ssim_val

    # 计算平均 SSIM 值
    ssim_avg = ssim_sum / len(filenames1)

    return ssim_avg

# 调用函数计算两个文件夹下所有图片的 SSIM 平均值
folder1 = "folder1"
folder2 = "folder2"
ssim_avg = calc_ssim_folder(folder1, folder2)

if ssim_avg is not None:
    print("两个文件夹下所有图片的 SSIM 平均值为:", ssim_avg)