# Source code for kornia.losses.psnr

import torch
import torch.nn as nn
from torch.nn.functional import mse_loss as mse

[docs]def psnr(input: torch.Tensor, target: torch.Tensor, max_val: float) -> torch.Tensor:
r"""Creates a function that calculates the PSNR between 2 images.

PSNR is Peek Signal to Noise Ratio, which is similar to mean squared error.
Given an m x n image, the PSNR is:

.. math::

\text{PSNR} = 10 \log_{10} \bigg(\frac{\text{MAX}_I^2}{MSE(I,T)}\bigg)

where

.. math::

\text{MSE}(I,T) = \frac{1}{mn}\sum_{i=0}^{m-1}\sum_{j=0}^{n-1} [I(i,j) - T(i,j)]^2

and :math:\text{MAX}_I is the maximum possible input value
(e.g for floating point images :math:\text{MAX}_I=1).

Args:
input (torch.Tensor): the input image with arbitrary shape :math:(*).
labels (torch.Tensor): the labels image with arbitrary shape :math:(*).
max_val (float): The maximum value in the input tensor.

Return:
torch.Tensor: the computed loss as a scalar.

Examples:
>>> ones = torch.ones(1)
>>> psnr(ones, 1.2 * ones, 2.) # 10 * log(4/((1.2-1)**2)) / log(10)
tensor(20.0000)

Reference:
https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio#Definition
"""
if not isinstance(input, torch.Tensor):
raise TypeError(f"Expected torch.Tensor but got {type(target)}.")

if not isinstance(target, torch.Tensor):
raise TypeError(f"Expected torch.Tensor but got {type(input)}.")

if input.shape != target.shape:
raise TypeError(f"Expected tensors of equal shapes, but got {input.shape} and {target.shape}")

return 10. * torch.log10(max_val ** 2 / mse(input, target, reduction='mean'))

[docs]def psnr_loss(input: torch.Tensor, target: torch.Tensor, max_val: float) -> torch.Tensor:
r"""Function that computes the PSNR loss.

The loss is computed as follows:

.. math::

\text{loss} = -\text{psnr(x, y)}

See :meth:~kornia.losses.psnr for details abut PSNR.

Args:
input (torch.Tensor): the input image with shape :math:(*).
labels (torch.Tensor): the labels image with shape :math:(*).
max_val (float): The maximum value in the input tensor.

Return:
torch.Tensor: the computed loss as a scalar.

Examples:
>>> ones = torch.ones(1)
>>> psnr_loss(ones, 1.2 * ones, 2.) # 10 * log(4/((1.2-1)**2)) / log(10)
tensor(-20.0000)
"""

return -1. * psnr(input, target, max_val)

[docs]class PSNRLoss(nn.Module):
r"""Creates a criterion that calculates the PSNR loss.

The loss is computed as follows:

.. math::

\text{loss} = -\text{psnr(x, y)}

See :meth:~kornia.losses.psnr for details abut PSNR.

Shape:
- Input: arbitrary dimensional tensor :math:(*).
- Target: arbitrary dimensional tensor :math:(*) same shape as input.
- Output: a scalar.

Examples:
>>> ones = torch.ones(1)
>>> criterion = PSNRLoss(2.)
>>> criterion(ones, 1.2 * ones) # 10 * log(4/((1.2-1)**2)) / log(10)
tensor(-20.0000)
"""

def __init__(self, max_val: float) -> None:
super(PSNRLoss, self).__init__()
self.max_val: float = max_val

def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
return psnr_loss(input, target, self.max_val)