Source code for viqa.fr_metrics.psnr
"""Module for calculating the peak signal-to-noise ratio (PSNR) between two images.
Examples
--------
.. doctest-requires:: numpy
>>> import numpy as np
>>> from viqa import PSNR
>>> img_r = np.zeros((256, 256))
>>> img_m = np.ones((256, 256))
>>> psnr = PSNR(data_range=1, normalize=False)
>>> psnr
PSNR(score_val=None)
>>> score = psnr(img_m, img_r)
>>> score
0.0
>>> psnr.print_score()
PSNR: 0.0
>>> img_r = np.zeros((256, 256))
>>> img_m = np.zeros((256, 256))
>>> psnr(img_m, img_r)
inf
>>> img_r = np.random.rand(256, 256)
>>> img_m = np.random.rand(128, 128)
>>> psnr(img_m, img_r)
Traceback (most recent call last):
...
ValueError: Image shapes do not match
"""
# Authors
# -------
# Author: Lukas Behammer
# Research Center Wels
# University of Applied Sciences Upper Austria, 2023
# CT Research Group
#
# Modifications
# -------------
# Original code, 2024, Lukas Behammer
#
# License
# -------
# BSD-3-Clause License
from warnings import warn
import numpy as np
from skimage.metrics import peak_signal_noise_ratio
from viqa._metrics import FullReferenceMetricsInterface
[docs]
class PSNR(FullReferenceMetricsInterface):
"""Class to calculate the peak signal-to-noise ratio (PSNR) between two images.
Attributes
----------
score_val : float
PSNR score value of the last calculation.
parameters : dict
Dictionary containing the parameters for PSNR calculation.
Parameters
----------
data_range : {1, 255, 65535}, default=255
Data range of the returned data in data loading. Is used for image loading when
``normalize`` is True and for the PSNR calculation. Passed to
:py:func:`viqa.utils.load_data` and :py:meth:`__call__`.
normalize : bool, default False
If True, the input images are normalized to the ``data_range`` argument.
**kwargs : optional
Additional parameters for data loading. The keyword arguments are passed to
:py:func:`viqa.utils.load_data`.
Other Parameters
----------------
chromatic : bool, default False
If True, the input images are expected to be RGB images.
If False, the input images are converted to grayscale images if necessary.
Raises
------
ValueError
If ``data_range`` is not set.
Notes
-----
``data_range`` for image loading is also used for the PSNR calculation and therefore
must be set. The parameter is set through the constructor of the class and is passed
to :py:meth:`__call__`.
"""
def __init__(self, data_range=255, normalize=False, **kwargs) -> None:
"""Construct method."""
if data_range is None:
raise ValueError("Parameter data_range must be set.")
super().__init__(data_range=data_range, normalize=normalize, **kwargs)
if self.parameters["chromatic"]:
self._name = "PSNRc"
else:
self._name = "PSNR"
def __call__(self, img_m, img_r):
"""Calculate the peak signal-to-noise ratio (PSNR) between two images.
Parameters
----------
img_m : np.ndarray or Tensor or str or os.PathLike
Distorted image to calculate score of.
img_r : np.ndarray or Tensor or str or os.PathLike
Reference image to calculate score against.
Returns
-------
score_val : float
PSNR score value.
"""
img_r, img_m = self.load_images(img_r, img_m)
# Calculate score
if np.array_equal(img_r, img_m):
score_val = np.inf # PSNR of identical images is infinity
else:
score_val = peak_signal_noise_ratio(
img_r, img_m, data_range=self.parameters["data_range"]
)
self.score_val = score_val
return score_val
[docs]
def print_score(self, decimals=2):
"""Print the PSNR score value of the last calculation.
Parameters
----------
decimals : int, default=2
Number of decimal places to print the score value.
Warns
-----
RuntimeWarning
If :py:attr:`score_val` is not available.
"""
if self.score_val is not None:
print("PSNR: {}".format(np.round(self.score_val, decimals)))
else:
warn("No score value for PSNR. Calculate the score first.", RuntimeWarning)