超分辨率指标计算：Python代码用于评估图像质量，PSNR、SSIM、LPIPS和NIQE

整体目标与设计思想

整体目标

这个Python脚本主要用于评估图像质量，它比较了一组高清（HD）图像和对应的生成图像，并计算了四种不同的图像质量指标：PSNR、SSIM、LPIPS和NIQE。

设计思想

1. 模块化: 代码通过函数进行模块化，每个函数负责一个特定任务，如计算PSNR或SSIM。
2. 并行化: 使用多进程来加速图像处理，特别是在处理大量图像时。
3. 可扩展性: 可以容易地添加更多的图像质量指标。
4. 结果记录: 为每个高清图像生成一个详细的文本文件，记录与之相关的生成图像的所有质量指标。

功能模块解析

初始化LPIPS模型

在代码开始，使用了LPIPS库来初始化一个预训练的VGG网络，这用于后续的LPIPS图像质量评估。

loss_fn = lpips.LPIPS(net='vgg')

图像质量评估函数

calculate_psnr, calculate_ssim

这两个函数分别使用OpenCV和skimage库来计算PSNR和SSIM。这些都是全参考指标，需要原图和生成图进行比较。

calculate_lpips

这个函数使用初始化的LPIPS模型来评估两个图像（原图和生成图）之间的感知差异。

calculate_niqe

这个函数实现了NIQE（无参考图像质量评估），即只需要一个图像即可评估其质量。

主要的图像处理函数：process_image()

这个函数是代码的核心，它执行以下操作：

1. 读取高清图像和生成图像: 使用OpenCV的cv2.imread方法。
2. 初始化质量指标: 初始化用于存储最佳质量指标的变量。
3. 计算并记录指标: 对于每个生成图像，计算所有的质量指标，并记录最佳值。
4. 结果存储: 将计算的质量指标保存到一个文本文件。

并行处理

使用Python的multiprocessing.Pool来并行处理所有高清图像。这是一种典型的"Map-Reduce"模式，其中process_image函数是map操作。

with Pool(4) as pool:  # Initialize a pool with 4 processes
    pool.starmap(process_image, [(i, main_output_file_path, hd_img_folder, generated_img_root_folder, output_root_folder) for i in range(1, 570)])

性能和优化

1. IO操作: 由于代码涉及大量的文件读写（图像和文本），IO可能成为性能瓶颈。
2. 并行优化: 使用多进程可以显著加速处理时间，但要注意不要超过系统的CPU核心数。

最佳实践与改进建议

1. 配置文件或命令行参数: 硬编码（如路径和进程数）应避免，最好使用配置文件或命令行参数。
2. 错误处理: 当前代码没有太多的错误处理，建议添加更多的异常捕获和错误检查。

代码

import os
from multiprocessing import Pool

import cv2
import lpips
import numpy as np
import torch
from scipy.ndimage import filters
from scipy.special import gammaln
from scipy.stats import genpareto
from skimage import img_as_float
from skimage.metrics import structural_similarity as compare_ssim
from tqdm import tqdm

# Initialize LPIPS
loss_fn = lpips.LPIPS(net='vgg')


def calculate_psnr(img1, img2):
    return cv2.PSNR(img1, img2)


def calculate_ssim(img1, img2):
    return compare_ssim(img1, img2, multichannel=True)


def calculate_lpips(img1, img2):
    img1 = torch.Tensor(img1).permute(2, 0, 1).unsqueeze(0)
    img2 = torch.Tensor(img2).permute(2, 0, 1).unsqueeze(0)
    return loss_fn(img1, img2).item()


def calculate_niqe(image):
    image = img_as_float(image)
    h, w = image.shape[:2]
    block_size = 96
    strides = 32
    features = []

    for i in range(0, h - block_size + 1, strides):
        for j in range(0, w - block_size + 1, strides):
            block = image[i:i + block_size, j:j + block_size]
            mu = np.mean(block)
            sigma = np.std(block)
            filtered_block = filters.gaussian_filter(block, sigma)
            shape, _, scale = genpareto.fit(filtered_block.ravel(), floc=0)
            feature = [mu, sigma, shape, scale, gammaln(1 / shape)]
            features.append(feature)

    features = np.array(features)
    model_mean = np.zeros(features.shape[1])
    model_cov_inv = np.eye(features.shape[1])
    quality_scores = []

    for feature in features:
        score = (feature - model_mean) @ model_cov_inv @ (feature - model_mean).T
        quality_scores.append(score)

    return np.mean(quality_scores)


def process_image(i, main_output_file_path, hd_img_folder, generated_img_root_folder, output_root_folder):
    hd_img_name = f"{i}.png"
    hd_img_path = os.path.join(hd_img_folder, hd_img_name)
    hd_img = cv2.imread(hd_img_path)

    corresponding_generated_folder = os.path.join(generated_img_root_folder, str(i))
    if not os.path.exists(corresponding_generated_folder):
        print(f"Folder for {hd_img_name} does not exist. Skipping.")
        return

    output_file_path = os.path.join(output_root_folder, f"{i}_output.txt")
    generated_img_names = os.listdir(corresponding_generated_folder)
    generated_img_names.sort(key=lambda x: int(x.split('.')[0]))
    total_images = len(generated_img_names)

    best_psnr = 0
    best_ssim = 0
    best_lpips = float('inf')
    best_niqe = float('inf')
    best_metrics_record = {}

    with open(output_file_path, 'w') as f:
        f.write(f"Results for HD Image: {hd_img_name}\n")
        f.write("-------------------------------------\n")

        for generated_img_name in tqdm(generated_img_names, total=total_images, desc=f"Processing for {hd_img_name}",
                                       leave=False):
            generated_img_path = os.path.join(corresponding_generated_folder, generated_img_name)
            generated_img = cv2.imread(generated_img_path)

            psnr = calculate_psnr(hd_img, generated_img)
            ssim = calculate_ssim(hd_img, generated_img)
            lpips_value = calculate_lpips(hd_img, generated_img)
            niqe = calculate_niqe(generated_img)

            result_line = f"{generated_img_name} PSNR: {psnr} SSIM: {ssim} LPIPS: {lpips_value} NIQE: {niqe}\n"
            f.write(result_line)

            if psnr > best_psnr:
                best_psnr = psnr
                best_metrics_record['Best PSNR'] = (generated_img_name, best_psnr)
            if ssim > best_ssim:
                best_ssim = ssim
                best_metrics_record['Best SSIM'] = (generated_img_name, best_ssim)
            if lpips_value < best_lpips:
                best_lpips = lpips_value
                best_metrics_record['Best LPIPS'] = (generated_img_name, best_lpips)
            if niqe < best_niqe:
                best_niqe = niqe
                best_metrics_record['Best NIQE'] = (generated_img_name, best_niqe)

        with open(main_output_file_path, 'a') as main_f:
            main_f.write(f"Best Metrics for {hd_img_name}\n")
            main_f.write("-------------------------------------\n")
            for metric, (img_name, value) in best_metrics_record.items():
                main_f.write(f"{metric}: {img_name}, Value: {value}\n")
            main_f.write("\n")
            print(f"Best Metrics for {hd_img_name} are saved in {main_output_file_path}")


if __name__ == "__main__":
    hd_img_folder = 'xxxxxxx'
    generated_img_root_folder = 'xxxxxxx'
    output_root_folder = 'xxxxxxx'

    main_output_file_path = os.path.join(output_root_folder, "all_results.txt")

    with open(main_output_file_path, 'w') as main_f:
        main_f.write("Detailed Results for Each High-Definition Image\n")
        main_f.write("==============================================\n")

    with Pool(4) as pool:  # Initialize a pool with 8 processes
        pool.starmap(process_image,
                     [(i, main_output_file_path, hd_img_folder, generated_img_root_folder, output_root_folder) for i in
                      range(1, 570)])  # Parallel processing