Source code for kornia.losses.psnr
import torch
import torch.nn as nn
import kornia.metrics as metrics
[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: the input image with shape :math:`(*)`.
labels : the labels image with shape :math:`(*)`.
max_val: The maximum value in the input tensor.
Return:
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.0 * metrics.psnr(input, target, max_val)
[docs]class PSNRLoss(nn.Module):
r"""Create 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.
Args:
max_val: The maximum value in the input tensor.
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().__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)