Psnr对比图片画质

 1import numpy as np
 2import cv2
 3import matplotlib.pyplot as plt
 4
 5
 6"""计算两幅RGB图像的PSNR值
 7参数:
 8    img1: 原始图像(numpy数组)
 9    img2: 待评估图像(numpy数组)
10返回:
11    PSNR值(dB)
12"""
13def rgb_psnr(img1, img2):
14    # 确保输入是numpy数组
15    img1 = np.array(img1, dtype=np.float64)
16    img2 = np.array(img2, dtype=np.float64)
17    # 计算各通道MSE
18    mse_r = np.mean((img1[:, :, 0] - img2[:, :, 0]) ** 2)
19    mse_g = np.mean((img1[:, :, 1] - img2[:, :, 1]) ** 2)
20    mse_b = np.mean((img1[:, :, 2] - img2[:, :, 2]) ** 2)
21    # 计算平均MSE
22    mse = (mse_r + mse_g + mse_b) / 3
23    # 处理完全相同的图像
24    if mse == 0:
25        return float('inf')
26    return 20 * np.log10(255 / np.sqrt(mse))
27
28
29def psnr(img1, img2):
30    img1 = img1.astype(np.float64)
31    img2 = img2.astype(np.float64)
32
33    mse = np.mean((img1 - img2) ** 2)
34    if mse == 0:
35        return float('inf')
36
37    if img1.max() > 1:
38        max_pixel = 255.0
39    else:
40        max_pixel = 1.0
41
42    return 20 * np.log10(max_pixel / np.sqrt(mse))
43
44
45if __name__ == "__main__":
46    # 1. 读取原始图像
47    original = cv2.imread('../../mario.png')  # BGR格式
48    original = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)  # 转换为RGB
49
50    # 2. 创建测试图像（添加高斯噪声）
51    noisy = original + np.random.normal(0, 25, original.shape)
52    noisy = np.clip(noisy, 0, 255).astype(np.uint8)
53
54    # 3. 计算PSNR
55    # psnr_value = rgb_psnr(original, noisy)
56    psnr_value = psnr(original, noisy)
57    print(f"PSNR between original and noisy image: {psnr_value:.2f} dB")
58
59    # 4. 可视化比较
60    plt.figure(figsize=(12, 6))
61
62    plt.subplot(1, 2, 1)
63    plt.imshow(original)
64    plt.title('Original Image')
65    plt.axis('off')
66
67    plt.subplot(1, 2, 2)
68    plt.imshow(noisy)
69    plt.title(f'Noisy Image (PSNR={psnr_value:.2f}dB)')
70    plt.axis('off')
71
72    plt.tight_layout()
73    plt.show()