TheAlgorithms-Python
46 строк · 1.3 Кб
1"""
2Peak signal-to-noise ratio - PSNR
3https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
4Source:
5https://tutorials.techonical.com/how-to-calculate-psnr-value-of-two-images-using-python
6"""
7
8import math
9import os
10
11import cv2
12import numpy as np
13
14PIXEL_MAX = 255.0
15
16
17def peak_signal_to_noise_ratio(original: float, contrast: float) -> float:
18mse = np.mean((original - contrast) ** 2)
19if mse == 0:
20return 100
21
22return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
23
24
25def main() -> None:
26dir_path = os.path.dirname(os.path.realpath(__file__))
27# Loading images (original image and compressed image)
28original = cv2.imread(os.path.join(dir_path, "image_data/original_image.png"))
29contrast = cv2.imread(os.path.join(dir_path, "image_data/compressed_image.png"), 1)
30
31original2 = cv2.imread(os.path.join(dir_path, "image_data/PSNR-example-base.png"))
32contrast2 = cv2.imread(
33os.path.join(dir_path, "image_data/PSNR-example-comp-10.jpg"), 1
34)
35
36# Value expected: 29.73dB
37print("-- First Test --")
38print(f"PSNR value is {peak_signal_to_noise_ratio(original, contrast)} dB")
39
40# # Value expected: 31.53dB (Wikipedia Example)
41print("\n-- Second Test --")
42print(f"PSNR value is {peak_signal_to_noise_ratio(original2, contrast2)} dB")
43
44
45if __name__ == "__main__":
46main()
47