cal.py

import glob
import os
import cv2
import numpy as np
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import peak_signal_noise_ratio as psnr
from tqdm import tqdm
import argparse


def load_image(image_path):
    """加载图片并转换为RGB格式"""
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError(f"无法加载图片: {image_path}")
    # OpenCV默认是BGR，转换为RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image


def calculate_psnr(img1, img2):
    """计算两张图片的PSNR"""
    # 确保图片尺寸相同
    if img1.shape != img2.shape:
        # 将图片调整为相同尺寸（使用较小的尺寸）
        h = min(img1.shape[0], img2.shape[0])
        w = min(img1.shape[1], img2.shape[1])
        img1 = img1[:h, :w]
        img2 = img2[:h, :w]

    return psnr(img1, img2, data_range=255)


def calculate_ssim(img1, img2):
    """计算两张图片的SSIM"""
    # 确保图片尺寸相同
    if img1.shape != img2.shape:
        h = min(img1.shape[0], img2.shape[0])
        w = min(img1.shape[1], img2.shape[1])
        img1 = img1[:h, :w]
        img2 = img2[:h, :w]

    # 如果是彩色图片，需要指定channel_axis
    if len(img1.shape) == 3:
        return ssim(img1, img2, channel_axis=2, data_range=255)
    else:
        return ssim(img1, img2, data_range=255)


def get_image_files(folder_path):
    """获取文件夹中的所有图片文件，按文件名排序"""
    # 支持的图片格式
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}

    image_files = []

    # 获取文件夹中的所有图片文件
    for file in os.listdir(folder_path):
        if os.path.splitext(file.lower())[1] in image_extensions:
            image_files.append(file)

    # 按文件名排序
    image_files.sort()

    return image_files


def compare_images(restored_folder, original_folder):
    """按顺序比较两个文件夹中的图片"""

    # 检查文件夹是否存在
    if not os.path.exists(restored_folder):
        raise ValueError(f"修复图片文件夹不存在: {restored_folder}")
    if not os.path.exists(original_folder):
        raise ValueError(f"原图文件夹不存在: {original_folder}")

    # 获取两个文件夹中的图片文件列表（按文件名排序）
    restored_files = get_image_files(restored_folder)
    original_files = get_image_files(original_folder)

    if not restored_files:
        print(f"修复文件夹 {restored_folder} 中没有找到图片文件！")
        return

    if not original_files:
        print(f"原图文件夹 {original_folder} 中没有找到图片文件！")
        return

    # 确定要处理的图片数量（取较小的数量）
    num_images = min(len(restored_files), len(original_files))

    if len(restored_files) != len(original_files):
        print(f"警告: 修复文件夹有 {len(restored_files)} 张图片，原图文件夹有 {len(original_files)} 张图片")
        print(f"将按顺序比较前 {num_images} 张图片")

    print(f"将按顺序比较 {num_images} 对图片")

    psnr_values = []
    ssim_values = []
    failed_pairs = []

    # 使用tqdm显示进度条
    for i in tqdm(range(num_images), desc="处理图片", unit="张"):
        try:
            # 按顺序获取对应的图片文件
            restored_file = restored_files[i]
            original_file = original_files[i]

            # 构建完整路径
            restored_path = os.path.join(restored_folder, restored_file)
            original_path = os.path.join(original_folder, original_file)

            # 加载图片
            restored_img = load_image(restored_path)
            original_img = load_image(original_path)

            # 计算PSNR和SSIM
            psnr_value = calculate_psnr(original_img, restored_img)
            ssim_value = calculate_ssim(original_img, restored_img)

            psnr_values.append(psnr_value)
            ssim_values.append(ssim_value)

        except Exception as e:
            failed_pairs.append((i + 1, restored_files[i], original_files[i], str(e)))
            print(f"\n处理第 {i + 1} 对图片时出错: {restored_files[i]} vs {original_files[i]} - {e}")

    # 计算统计结果
    if psnr_values and ssim_values:
        avg_psnr = np.mean(psnr_values)
        avg_ssim = np.mean(ssim_values)
        std_psnr = np.std(psnr_values)
        std_ssim = np.std(ssim_values)
        min_psnr = np.min(psnr_values)
        max_psnr = np.max(psnr_values)
        min_ssim = np.min(ssim_values)
        max_ssim = np.max(ssim_values)

        print("\n" + "=" * 60)
        print(f"文件夹: {os.path.basename(restored_folder)} 的图片质量评估结果")
        print("=" * 60)
        print(f"成功处理的图片数量: {len(psnr_values)}")
        print(f"平均 PSNR: {avg_psnr:.4f} dB")
        print(f"平均 SSIM: {avg_ssim:.4f}")
        print(f"PSNR 标准差: {std_psnr:.4f}")
        print(f"SSIM 标准差: {std_ssim:.4f}")
        print(f"PSNR 范围: {min_psnr:.4f} - {max_psnr:.4f}")
        print(f"SSIM 范围: {min_ssim:.4f} - {max_ssim:.4f}")

        if failed_pairs:
            print(f"\n失败的图片对数量: {len(failed_pairs)}")
            for pair_num, restored_file, original_file, error in failed_pairs:
                print(f"  - 第{pair_num}对: {restored_file} vs {original_file} - {error}")

        return {
            'folder_name': os.path.basename(restored_folder),
            'image_count': len(psnr_values),
            'avg_psnr': avg_psnr,
            'avg_ssim': avg_ssim,
            'std_psnr': std_psnr,
            'std_ssim': std_ssim,
            'min_psnr': min_psnr,
            'max_psnr': max_psnr,
            'min_ssim': min_ssim,
            'max_ssim': max_ssim,
            'failed_count': len(failed_pairs)
        }
    else:
        print("没有成功处理任何图片！")
        return None


def print_summary(all_results):
    """打印所有结果的汇总"""
    if not all_results:
        print("没有成功处理任何文件夹！")
        return

    print("\n" + "=" * 80)
    print("所有文件夹结果汇总")
    print("=" * 80)

    # 打印对比表格
    print(f"{'文件夹名':<20} {'图片数':<8} {'平均PSNR':<12} {'平均SSIM':<12} {'失败数':<8}")
    print("-" * 80)

    for result in all_results:
        print(f"{result['folder_name']:<20} {result['image_count']:<8} "
              f"{result['avg_psnr']:<12.4f} {result['avg_ssim']:<12.4f} "
              f"{result['failed_count']:<8}")

    # 找出最佳结果
    print("\n" + "=" * 80)
    print("最佳结果")
    print("=" * 80)

    best_psnr_result = max(all_results, key=lambda x: x['avg_psnr'])
    best_ssim_result = max(all_results, key=lambda x: x['avg_ssim'])

    print(f"🏆 最佳 PSNR: {best_psnr_result['folder_name']} ({best_psnr_result['avg_psnr']:.4f} dB)")
    print(f"🏆 最佳 SSIM: {best_ssim_result['folder_name']} ({best_ssim_result['avg_ssim']:.4f})")


if __name__ == "__main__":
    restored_folder_prefix = "/opt/data/private/dcc/transformer/results_celeba*/restored"  # 修复后图片文件夹的匹配模式
    original_folder = "/opt/data/private/celeba256_test"  # 原图文件夹

    # 查找所有匹配的文件夹
    folders = glob.glob(restored_folder_prefix)
    folders = [f for f in folders if os.path.isdir(f)]
    folders.sort()

    if not folders:
        print(f"没有找到匹配的文件夹: {restored_folder_prefix}")
    else:
        print(f"找到 {len(folders)} 个文件夹:")
        for folder in folders:
            print(f"  - {os.path.basename(folder)}")
        print()

        all_results = []

        # 使用外层进度条显示文件夹处理进度
        for folder in tqdm(folders, desc="处理文件夹", unit="个"):
            print(f"\n正在处理文件夹: {os.path.basename(folder)}")
            result = compare_images(folder, original_folder)
            if result:
                all_results.append(result)

        # 打印汇总结果
        print_summary(all_results)