test.py

# -*- coding: utf-8 -*-
# @Author:
# @Date: 2025-06-11 15:25:37
# @LastEditTime: 2025-06-28 18:30:00
# @FilePath: /image_denoise/main.py
# @Description: 主函数，执行图像去噪与全面的性能评估
import time
from enum import Enum
import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image

# 导入 skimage 用于计算 PSNR 和 SSIM
from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.metrics import structural_similarity as ssim

from Method import Method
from TotalVariation import TV_L1, TV_L2
from pnp import PnP_BM3D, PnP_CNN


# ... (load_img, noise_img, save_img, Methods Enum - 这些函数保持不变) ...
def load_img(path: str):
    """
    加载图像文件
    """
    return Image.open(path)


def noise_img(img: np.ndarray, noise_sigma: float):
    """
    添加高斯噪声到图像
    """
    temp_image = np.float64(np.copy(img))
    np.random.seed(42)
    noise = np.random.normal(0, noise_sigma, temp_image.shape) * 255
    noisy_image = temp_image + noise
    return noisy_image


def save_img(img: np.ndarray, path: str):
    """
    保存图像文件
    """
    os.makedirs(os.path.dirname(path), exist_ok=True)
    arr = np.clip(img, 0, 255).astype(np.uint8)
    if arr.ndim == 3 and arr.shape[2] == 1:
        arr = arr.squeeze(axis=2)
    img_pil = Image.fromarray(arr)
    img_pil.save(path)


class Methods(Enum):
    TV_L1 = "TV-L1"
    TV_L2 = "TV-L2"
    PnP_BM3D = "PnP-BM3D"
    PnP_CNN = "PnP-CNN"


def run_single_test(method_enum: Methods, params: dict, original_img: np.ndarray, noisy_img: np.ndarray):
    """
    运行单次去噪测试并返回包含性能指标的字典。
    """
    method_name = method_enum.value
    print(f"--- Running: {method_name} with params: {params} ---")

    # 根据方法枚举选择对应的类
    method_class_map = {
        Methods.TV_L1: TV_L1,
        Methods.TV_L2: TV_L2,
        Methods.PnP_BM3D: PnP_BM3D,
        Methods.PnP_CNN: PnP_CNN,
    }
    method_class = method_class_map.get(method_enum)
    if not method_class:
        raise NotImplementedError(f"Method {method_name} is not implemented in the runner.")

    denoised_method = method_class(noisy_img, **params)

    start_time = time.time()
    denoised_img = denoised_method.solve()
    execution_time = time.time() - start_time

    # 计算评估指标
    final_psnr = psnr(original_img, denoised_img, data_range=255)
    final_ssim = ssim(original_img, denoised_img, data_range=255)

    print(f"Result: PSNR={final_psnr:.2f} dB, SSIM={final_ssim:.4f}, Time={execution_time:.2f}s")

    return {
        'method_enum': method_enum,
        'method_name': method_name,
        'params': params,
        'psnr': final_psnr,
        'ssim': final_ssim,
        'time': execution_time,
        'denoised_img': denoised_img
    }


def plot_parameter_sweep_charts(results: list, title: str, variable_param: str, output_dir: str):
    """
    为参数扫描实验绘制并保存对比图表。
    """
    if not results: return

    param_values = [str(r['params'][variable_param]) for r in results]
    psnr_values = [r['psnr'] for r in results]
    ssim_values = [r['ssim'] for r in results]
    time_values = [r['time'] for r in results]

    fig, axs = plt.subplots(1, 3, figsize=(20, 5))
    fig.suptitle(title, fontsize=14)

    # PSNR Chart
    axs[0].bar(param_values, psnr_values, color='skyblue')
    axs[0].set_title('PSNR (Higher is Better)')
    axs[0].set_xlabel(variable_param)
    axs[0].set_ylabel('PSNR (dB)')
    for i, v in enumerate(psnr_values): axs[0].text(i, v, f"{v:.2f}", ha='center', va='bottom')

    # SSIM Chart
    axs[1].bar(param_values, ssim_values, color='lightgreen')
    axs[1].set_title('SSIM (Higher is Better)')
    axs[1].set_xlabel(variable_param)
    axs[1].set_ylabel('SSIM')
    axs[1].set_ylim(min(ssim_values) - 0.005, max(ssim_values) + 0.005)
    for i, v in enumerate(ssim_values): axs[1].text(i, v, f"{v:.4f}", ha='center', va='bottom')

    # Time Chart
    axs[2].bar(param_values, time_values, color='salmon')
    axs[2].set_title('Execution Time (Lower is Better)')
    axs[2].set_xlabel(variable_param)
    axs[2].set_ylabel('Time (s)')
    for i, v in enumerate(time_values): axs[2].text(i, v, f"{v:.2f}", ha='center', va='bottom')

    plt.tight_layout(rect=[0, 0, 1, 0.95])
    filename = title.replace(" ", "_").replace("=", "").replace(",", "") + ".png"
    plt.savefig(os.path.join(output_dir, filename))
    plt.close()


def test():
    """
    执行全面的、控制变量的实验，评估和比较去噪方法。
    """
    # 1. 准备图像
    try:
        ori_img_pil = load_img('image.png')
    except FileNotFoundError:
        print("Warning: 'image.png' not found. Creating a dummy image.")
        ori_img_pil = Image.fromarray(np.zeros((256, 256), dtype=np.uint8) + 128)

    gray_img = np.asarray(ori_img_pil.convert('L'), np.float64)
    noisy_img = noise_img(gray_img, noise_sigma=0.1)

    output_dir = "test_results_comprehensive"
    os.makedirs(output_dir, exist_ok=True)
    print(f"Saving all results to '{output_dir}/' directory.")

    # 2. 定义所有方法的参数扫描配置
    sweep_configs = {
        Methods.TV_L1: {
            'base_params': {'max_iter': 500, 'cost_threshold': 1e-3},
            'lambda_sweep': {'rho': 0.07, 'lambda_': [0.5, 1.0, 1.5, 2.0]},
            'rho_sweep': {'lambda_': 1.0, 'rho': [0.03, 0.05, 0.07, 0.1]}
        },
        Methods.TV_L2: {
            'base_params': {'max_iter': 500, 'cost_threshold': 1e-3},
            'lambda_sweep': {'rho': 0.07, 'lambda_': [4.0, 7.0, 10.0, 13.0]},
            'rho_sweep': {'lambda_': 7.0, 'rho': [0.03, 0.05, 0.07, 0.1]}
        },
        Methods.PnP_BM3D: {
            'base_params': {'max_iter': 30, 'cost_threshold': 1e-3},
            'lambda_sweep': {'rho': 0.03, 'lambda_': [0.5, 0.8, 1.0, 1.2]},
            'rho_sweep': {'lambda_': 1.0, 'rho': [0.01, 0.03, 0.05, 0.07]}
        },
        Methods.PnP_CNN: {
            'base_params': {'max_iter': 50, 'cost_threshold': 1e-3},
            'lambda_sweep': {'rho': 0.07, 'lambda_': [0.5, 0.7, 1.0, 1.2]},
            'rho_sweep': {'lambda_': 0.7, 'rho': [0.03, 0.05, 0.07, 0.1]}
        }
    }

    all_results = []
    best_results_for_final_comparison = []

    # 3. 执行所有参数扫描实验
    for method_enum, config in sweep_configs.items():
        method_name = method_enum.value
        print(f"\n{'=' * 20} Starting Parameter Sweep for {method_name} {'=' * 20}")

        method_all_runs = []

        # a) Lambda sweep
        lambda_sweep_params = config['lambda_sweep']
        rho_fixed = lambda_sweep_params['rho']
        lambda_results = []
        for l in lambda_sweep_params['lambda_']:
            params = {**config['base_params'], 'lambda_': l, 'rho': rho_fixed}
            result = run_single_test(method_enum, params, gray_img, noisy_img)
            lambda_results.append(result)
        plot_parameter_sweep_charts(
            lambda_results, f"{method_name} Performance vs lambda_ (rho={rho_fixed})", 'lambda_', output_dir
        )
        method_all_runs.extend(lambda_results)

        # b) Rho sweep
        rho_sweep_params = config['rho_sweep']
        lambda_fixed = rho_sweep_params['lambda_']
        rho_results = []
        for r in rho_sweep_params['rho']:
            params = {**config['base_params'], 'lambda_': lambda_fixed, 'rho': r}
            # 避免重复测试
            if not any(res['params'] == params for res in method_all_runs):
                result = run_single_test(method_enum, params, gray_img, noisy_img)
                rho_results.append(result)
        plot_parameter_sweep_charts(
            rho_results, f"{method_name} Performance vs rho (lambda_={lambda_fixed})", 'rho', output_dir
        )
        method_all_runs.extend(rho_results)

        # c) 找出当前方法的最佳结果 (基于SSIM)
        best_run_for_method = max(method_all_runs, key=lambda x: x['ssim'])
        best_results_for_final_comparison.append(best_run_for_method)
        all_results.extend(method_all_runs)
        print(
            f"--- Best parameters for {method_name} found: {best_run_for_method['params']} (SSIM: {best_run_for_method['ssim']:.4f}) ---")

    # 4. 执行最终的横向对比
    print(f"\n{'=' * 20} Starting Final Cross-Method Comparison {'=' * 20}")

    # 排序以便图表美观
    best_results_for_final_comparison.sort(key=lambda x: x['method_name'])

    # 生成图表标题，包含所有最优参数
    title_params = ' | '.join([f"{r['method_name']}: λ={r['params']['lambda_']}, ρ={r['params']['rho']}" for r in
                               best_results_for_final_comparison])

    method_names = [r['method_name'] for r in best_results_for_final_comparison]
    psnr_values = [r['psnr'] for r in best_results_for_final_comparison]
    ssim_values = [r['ssim'] for r in best_results_for_final_comparison]
    time_values = [r['time'] for r in best_results_for_final_comparison]

    fig, axs = plt.subplots(1, 3, figsize=(22, 6))
    fig.suptitle(f"Final Cross-Method Comparison (Using Best Params)\n{title_params}", fontsize=12)

    # PSNR Chart
    axs[0].bar(method_names, psnr_values, color='darkcyan')
    axs[0].set_title('PSNR (Higher is Better)')
    axs[0].set_ylabel('PSNR (dB)')
    for i, v in enumerate(psnr_values): axs[0].text(i, v / 2, f"{v:.2f}", ha='center', color='white', weight='bold')

    # SSIM Chart
    axs[1].bar(method_names, ssim_values, color='orchid')
    axs[1].set_title('SSIM (Higher is Better)')
    axs[1].set_ylabel('SSIM')
    for i, v in enumerate(ssim_values): axs[1].text(i, v / 2, f"{v:.4f}", ha='center', color='white', weight='bold')

    # Time Chart
    axs[2].bar(method_names, time_values, color='goldenrod')
    axs[2].set_title('Execution Time (Log Scale, Lower is Better)')
    axs[2].set_ylabel('Time (s)')
    axs[2].set_yscale('log')
    for i, v in enumerate(time_values): axs[2].text(i, v, f"{v:.2f}", ha='center', va='bottom')

    plt.tight_layout(rect=[0, 0.05, 1, 0.9])
    plt.savefig(os.path.join(output_dir, "Final_Cross_Method_Comparison.png"))
    plt.close()

    # 5. 保存所有去噪图像并生成最终报告
    print("\nGenerating final report and saving all denoised images...")
    report_path = os.path.join(output_dir, "final_summary_report.md")
    with open(report_path, 'w') as f:
        f.write("# Comprehensive Image Denoising Experiment Report\n\n")
        f.write("This report summarizes the performance of various denoising methods across different parameters.\n")
        f.write(
            "An asterisk (*) indicates the best performing parameter set for that method, used in the final comparison.\n\n")
        f.write("| Method       | Parameters (lambda, rho, max_iter) | PSNR (dB) | SSIM      | Time (s) | Best |\n")
        f.write("|--------------|------------------------------------|-----------|-----------|----------|------|\n")

        # 排序以方便查看
        all_results.sort(key=lambda x: (x['method_name'], x['params']['lambda_'], x['params']['rho']))

        for r in all_results:
            is_best = any(r['params'] == best_r['params'] and r['method_name'] == best_r['method_name'] for best_r in
                          best_results_for_final_comparison)
            star = ' *' if is_best else ''
            param_str = f"λ={r['params']['lambda_']}, ρ={r['params']['rho']}, iter={r['params']['max_iter']}"
            f.write(
                f"| {r['method_name']:<12} | {param_str:<34} | {r['psnr']:<9.2f} | {r['ssim']:<9.4f} | {r['time']:<8.2f} |{star}|\n")

            # 保存每个结果的去噪图像
            img_filename = f"{r['method_name']}_l-{r['params']['lambda_']}_r-{r['params']['rho']}.png"
            save_img(r['denoised_img'], os.path.join(output_dir, 'all_denoised_images', img_filename))

    print(f"\nAll tests completed. All charts, images, and a summary report are in '{output_dir}'.")


if __name__ == "__main__":
    test()