Psnr对比图片画质

import numpy as np
import cv2
import matplotlib.pyplot as plt


"""计算两幅RGB图像的PSNR值
参数:
    img1: 原始图像(numpy数组)
    img2: 待评估图像(numpy数组)
返回:
    PSNR值(dB)
"""
def rgb_psnr(img1, img2):
    # 确保输入是numpy数组
    img1 = np.array(img1, dtype=np.float64)
    img2 = np.array(img2, dtype=np.float64)
    # 计算各通道MSE
    mse_r = np.mean((img1[:, :, 0] - img2[:, :, 0]) ** 2)
    mse_g = np.mean((img1[:, :, 1] - img2[:, :, 1]) ** 2)
    mse_b = np.mean((img1[:, :, 2] - img2[:, :, 2]) ** 2)
    # 计算平均MSE
    mse = (mse_r + mse_g + mse_b) / 3
    # 处理完全相同的图像
    if mse == 0:
        return float('inf')
    return 20 * np.log10(255 / np.sqrt(mse))


def psnr(img1, img2):
    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)

    mse = np.mean((img1 - img2) ** 2)
    if mse == 0:
        return float('inf')

    if img1.max() > 1:
        max_pixel = 255.0
    else:
        max_pixel = 1.0

    return 20 * np.log10(max_pixel / np.sqrt(mse))


if __name__ == "__main__":
    # 1. 读取原始图像
    original = cv2.imread('../../mario.png')  # BGR格式
    original = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)  # 转换为RGB

    # 2. 创建测试图像（添加高斯噪声）
    noisy = original + np.random.normal(0, 25, original.shape)
    noisy = np.clip(noisy, 0, 255).astype(np.uint8)

    # 3. 计算PSNR
    # psnr_value = rgb_psnr(original, noisy)
    psnr_value = psnr(original, noisy)
    print(f"PSNR between original and noisy image: {psnr_value:.2f} dB")

    # 4. 可视化比较
    plt.figure(figsize=(12, 6))

    plt.subplot(1, 2, 1)
    plt.imshow(original)
    plt.title('Original Image')
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.imshow(noisy)
    plt.title(f'Noisy Image (PSNR={psnr_value:.2f}dB)')
    plt.axis('off')

    plt.tight_layout()
    plt.show()