Make data_range required parameter (#3178)

* change psnr and psnrb implementations
* fix tests
* changelog
* fix deprecated
* typing ?


---------

Co-authored-by: Jirka B <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -39,7 +39,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
--
+- Changed `data_range` argument in `PSNR` metric to be a required argument ([3178](https://github.com/Lightning-AI/torchmetrics/pull/3178))
 
 
 ### Deprecated

src/torchmetrics/functional/image/_deprecated.py

@@ -80,7 +80,7 @@ def _image_gradients(img: Tensor) -> tuple[Tensor, Tensor]:
 def _peak_signal_noise_ratio(
     preds: Tensor,
     target: Tensor,
-    data_range: Optional[Union[float, tuple[float, float]]] = None,
+    data_range: Union[float, tuple[float, float]] = 3.0,
     base: float = 10.0,
     reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
     dim: Optional[Union[int, tuple[int, ...]]] = None,

src/torchmetrics/functional/image/psnr.py

@@ -95,7 +95,7 @@ def _psnr_update(
 def peak_signal_noise_ratio(
     preds: Tensor,
     target: Tensor,
-    data_range: Optional[Union[float, tuple[float, float]]] = None,
+    data_range: Union[float, tuple[float, float]],
     base: float = 10.0,
     reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
     dim: Optional[Union[int, tuple[int, ...]]] = None,
@@ -106,9 +106,8 @@ def peak_signal_noise_ratio(
         preds: estimated signal
         target: groun truth signal
         data_range:
-            the range of the data. If None, it is determined from the data (max - min). If a tuple is provided then
-            the range is calculated as the difference and input is clamped between the values.
-            The ``data_range`` must be given when ``dim`` is not None.
+            the range of the data. If a tuple is provided then the range is calculated as the difference and
+            input is clamped between the values.
         base: a base of a logarithm to use
         reduction: a method to reduce metric score over labels.
 
@@ -123,15 +122,11 @@ def peak_signal_noise_ratio(
     Return:
         Tensor with PSNR score
 
-    Raises:
-        ValueError:
-            If ``dim`` is not ``None`` and ``data_range`` is not provided.
-
     Example:
         >>> from torchmetrics.functional.image import peak_signal_noise_ratio
         >>> pred = torch.tensor([[0.0, 1.0], [2.0, 3.0]])
         >>> target = torch.tensor([[3.0, 2.0], [1.0, 0.0]])
-        >>> peak_signal_noise_ratio(pred, target)
+        >>> peak_signal_noise_ratio(pred, target, data_range=3.0)
         tensor(2.5527)
 
     .. attention::
@@ -141,19 +136,12 @@ def peak_signal_noise_ratio(
     if dim is None and reduction != "elementwise_mean":
         rank_zero_warn(f"The `reduction={reduction}` will not have any effect when `dim` is None.")
 
-    if data_range is None:
-        if dim is not None:
-            # Maybe we could use `torch.amax(target, dim=dim) - torch.amin(target, dim=dim)` in PyTorch 1.7 to calculate
-            # `data_range` in the future.
-            raise ValueError("The `data_range` must be given when `dim` is not None.")
-
-        data_range = target.max() - target.min()  # type: ignore[assignment]
-    elif isinstance(data_range, tuple):
+    if isinstance(data_range, tuple):
         preds = torch.clamp(preds, min=data_range[0], max=data_range[1])
         target = torch.clamp(target, min=data_range[0], max=data_range[1])
-        data_range = tensor(data_range[1] - data_range[0])  # type: ignore[assignment]
+        data_range_val = tensor(data_range[1] - data_range[0])
     else:
-        data_range = tensor(float(data_range))  # type: ignore[assignment]
+        data_range_val = tensor(float(data_range))
 
     sum_squared_error, num_obs = _psnr_update(preds, target, dim=dim)
-    return _psnr_compute(sum_squared_error, num_obs, data_range, base=base, reduction=reduction)  # type: ignore[arg-type]
+    return _psnr_compute(sum_squared_error, num_obs, data_range_val, base=base, reduction=reduction)

src/torchmetrics/functional/image/psnrb.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
+from typing import Union
 
 import torch
 from torch import Tensor, tensor
@@ -76,13 +77,11 @@ def _psnrb_compute(
         sum_squared_error: Sum of square of errors over all observations
         bef: block effect
         num_obs: Number of predictions or observations
-        data_range: the range of the data. If None, it is determined from the data (max - min).
+        data_range: the range of the data.
 
     """
     sum_squared_error = sum_squared_error / num_obs + bef
-    if data_range > 2:
-        return 10 * torch.log10(data_range**2 / sum_squared_error)
-    return 10 * torch.log10(1.0 / sum_squared_error)
+    return 10 * torch.log10(data_range**2 / sum_squared_error)
 
 
 def _psnrb_update(preds: Tensor, target: Tensor, block_size: int = 8) -> tuple[Tensor, Tensor, Tensor]:
@@ -103,6 +102,7 @@ def _psnrb_update(preds: Tensor, target: Tensor, block_size: int = 8) -> tuple[T
 def peak_signal_noise_ratio_with_blocked_effect(
     preds: Tensor,
     target: Tensor,
+    data_range: Union[float, tuple[float, float]],
     block_size: int = 8,
 ) -> Tensor:
     r"""Computes `Peak Signal to Noise Ratio With Blocked Effect` (PSNRB) metrics.
@@ -114,7 +114,9 @@ def peak_signal_noise_ratio_with_blocked_effect(
 
     Args:
         preds: estimated signal
-        target: groun truth signal
+        target: ground truth signal
+        data_range: the range of the data. If a tuple is provided then the range is calculated as the difference and
+            input is clamped between the values.
         block_size: integer indication the block size
 
     Return:
@@ -125,10 +127,16 @@ def peak_signal_noise_ratio_with_blocked_effect(
         >>> from torchmetrics.functional.image import peak_signal_noise_ratio_with_blocked_effect
         >>> preds = rand(1, 1, 28, 28)
         >>> target = rand(1, 1, 28, 28)
-        >>> peak_signal_noise_ratio_with_blocked_effect(preds, target)
+        >>> peak_signal_noise_ratio_with_blocked_effect(preds, target, data_range=1.0)
         tensor(7.8402)
 
     """
-    data_range = target.max() - target.min()
+    if isinstance(data_range, tuple):
+        preds = torch.clamp(preds, min=data_range[0], max=data_range[1])
+        target = torch.clamp(target, min=data_range[0], max=data_range[1])
+        data_range_val = tensor(data_range[1] - data_range[0])
+    else:
+        data_range_val = tensor(float(data_range))
+
     sum_squared_error, bef, num_obs = _psnrb_update(preds, target, block_size=block_size)
-    return _psnrb_compute(sum_squared_error, bef, num_obs, data_range)
+    return _psnrb_compute(sum_squared_error, bef, num_obs, data_range_val)

src/torchmetrics/image/_deprecated.py

@@ -91,7 +91,7 @@ class _PeakSignalNoiseRatio(PeakSignalNoiseRatio):
 
     def __init__(
         self,
-        data_range: Optional[Union[float, tuple[float, float]]] = None,
+        data_range: Union[float, tuple[float, float]] = 3.0,
         base: float = 10.0,
         reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
         dim: Optional[Union[int, tuple[int, ...]]] = None,

src/torchmetrics/image/psnr.py

@@ -48,9 +48,8 @@ class PeakSignalNoiseRatio(Metric):
 
     Args:
         data_range:
-            the range of the data. If None, it is determined from the data (max - min). If a tuple is provided then
-            the range is calculated as the difference and input is clamped between the values.
-            The ``data_range`` must be given when ``dim`` is not None.
+            the range of the data. If a tuple is provided, then the range is calculated as the difference and
+            input is clamped between the values.
         base: a base of a logarithm to use.
         reduction: a method to reduce metric score over labels.
 
@@ -63,13 +62,9 @@ class PeakSignalNoiseRatio(Metric):
             None meaning scores will be reduced across all dimensions and all batches.
         kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
-    Raises:
-        ValueError:
-            If ``dim`` is not ``None`` and ``data_range`` is not given.
-
     Example:
         >>> from torchmetrics.image import PeakSignalNoiseRatio
-        >>> psnr = PeakSignalNoiseRatio()
+        >>> psnr = PeakSignalNoiseRatio(data_range=3.0)
         >>> preds = torch.tensor([[0.0, 1.0], [2.0, 3.0]])
         >>> target = torch.tensor([[3.0, 2.0], [1.0, 0.0]])
         >>> psnr(preds, target)
@@ -81,13 +76,11 @@ class PeakSignalNoiseRatio(Metric):
     higher_is_better: bool = True
     full_state_update: bool = False
     plot_lower_bound: float = 0.0
-
-    min_target: Tensor
-    max_target: Tensor
+    data_range: Tensor
 
     def __init__(
         self,
-        data_range: Optional[Union[float, tuple[float, float]]] = None,
+        data_range: Union[float, tuple[float, float]],
         base: float = 10.0,
         reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
         dim: Optional[Union[int, tuple[int, ...]]] = None,
@@ -106,20 +99,12 @@ def __init__(
             self.add_state("total", default=[], dist_reduce_fx="cat")
 
         self.clamping_fn = None
-        if data_range is None:
-            if dim is not None:
-                # Maybe we could use `torch.amax(target, dim=dim) - torch.amin(target, dim=dim)` in PyTorch 1.7 to
-                # calculate `data_range` in the future.
-                raise ValueError("The `data_range` must be given when `dim` is not None.")
-
-            self.data_range = None
-            self.add_state("min_target", default=tensor(0.0), dist_reduce_fx=torch.min)
-            self.add_state("max_target", default=tensor(0.0), dist_reduce_fx=torch.max)
-        elif isinstance(data_range, tuple):
+        if isinstance(data_range, tuple):
             self.add_state("data_range", default=tensor(data_range[1] - data_range[0]), dist_reduce_fx="mean")
             self.clamping_fn = partial(torch.clamp, min=data_range[0], max=data_range[1])
         else:
             self.add_state("data_range", default=tensor(float(data_range)), dist_reduce_fx="mean")
+
         self.base = base
         self.reduction = reduction
         self.dim = tuple(dim) if isinstance(dim, Sequence) else dim
@@ -132,11 +117,6 @@ def update(self, preds: Tensor, target: Tensor) -> None:
 
         sum_squared_error, num_obs = _psnr_update(preds, target, dim=self.dim)
         if self.dim is None:
-            if self.data_range is None:
-                # keep track of min and max target values
-                self.min_target = torch.minimum(target.min(), self.min_target)
-                self.max_target = torch.maximum(target.max(), self.max_target)
-
             if not isinstance(self.sum_squared_error, Tensor):
                 raise TypeError(
                     f"Expected `self.sum_squared_error` to be a Tensor, but got {type(self.sum_squared_error)}"
@@ -158,8 +138,6 @@ def update(self, preds: Tensor, target: Tensor) -> None:
 
     def compute(self) -> Tensor:
         """Compute peak signal-to-noise ratio over state."""
-        data_range = self.data_range if self.data_range is not None else self.max_target - self.min_target
-
         if isinstance(self.sum_squared_error, torch.Tensor):
             sum_squared_error = self.sum_squared_error
         elif isinstance(self.sum_squared_error, list):
@@ -174,7 +152,7 @@ def compute(self) -> Tensor:
         else:
             raise TypeError("Expected total to be a Tensor or a list of Tensors")
 
-        return _psnr_compute(sum_squared_error, total, data_range, base=self.base, reduction=self.reduction)
+        return _psnr_compute(sum_squared_error, total, self.data_range, base=self.base, reduction=self.reduction)
 
     def plot(
         self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
@@ -199,7 +177,7 @@ def plot(
             >>> # Example plotting a single value
             >>> import torch
             >>> from torchmetrics.image import PeakSignalNoiseRatio
-            >>> metric = PeakSignalNoiseRatio()
+            >>> metric = PeakSignalNoiseRatio(data_range=1.0)
             >>> preds = torch.tensor([[0.0, 1.0], [2.0, 3.0]])
             >>> target = torch.tensor([[3.0, 2.0], [1.0, 0.0]])
             >>> metric.update(preds, target)
@@ -211,7 +189,7 @@ def plot(
             >>> # Example plotting multiple values
             >>> import torch
             >>> from torchmetrics.image import PeakSignalNoiseRatio
-            >>> metric = PeakSignalNoiseRatio()
+            >>> metric = PeakSignalNoiseRatio(data_range=1.0)
             >>> preds = torch.tensor([[0.0, 1.0], [2.0, 3.0]])
             >>> target = torch.tensor([[3.0, 2.0], [1.0, 0.0]])
             >>> values = [ ]

src/torchmetrics/image/psnrb.py

@@ -48,12 +48,14 @@ class PeakSignalNoiseRatioWithBlockedEffect(Metric):
     - ``psnrb`` (:class:`~torch.Tensor`): float scalar tensor with aggregated PSNRB value
 
     Args:
+        data_range: the range of the data. If a tuple is provided then the range is calculated as the difference and
+            input is clamped between the values.
         block_size: integer indication the block size
         kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
     Example:
         >>> from torch import rand
-        >>> metric = PeakSignalNoiseRatioWithBlockedEffect()
+        >>> metric = PeakSignalNoiseRatioWithBlockedEffect(data_range=1.0)
         >>> preds = rand(2, 1, 10, 10)
         >>> target = rand(2, 1, 10, 10)
         >>> metric(preds, target)
@@ -72,6 +74,7 @@ class PeakSignalNoiseRatioWithBlockedEffect(Metric):
 
     def __init__(
         self,
+        data_range: Union[float, tuple[float, float]],
         block_size: int = 8,
         **kwargs: Any,
     ) -> None:
@@ -83,15 +86,24 @@ def __init__(
         self.add_state("sum_squared_error", default=tensor(0.0), dist_reduce_fx="sum")
         self.add_state("total", default=tensor(0), dist_reduce_fx="sum")
         self.add_state("bef", default=tensor(0.0), dist_reduce_fx="sum")
-        self.add_state("data_range", default=tensor(0), dist_reduce_fx="max")
+
+        if isinstance(data_range, tuple):
+            self.add_state("data_range", default=tensor(data_range[1] - data_range[0]), dist_reduce_fx="mean")
+            self.clamping_fn = lambda x: torch.clamp(x, min=data_range[0], max=data_range[1])
+        else:
+            self.add_state("data_range", default=tensor(float(data_range)), dist_reduce_fx="mean")
+            self.clamping_fn = None  # type: ignore[assignment]
 
     def update(self, preds: Tensor, target: Tensor) -> None:
         """Update state with predictions and targets."""
+        if self.clamping_fn is not None:
+            preds = self.clamping_fn(preds)
+            target = self.clamping_fn(target)
+
         sum_squared_error, bef, num_obs = _psnrb_update(preds, target, block_size=self.block_size)
         self.sum_squared_error += sum_squared_error
         self.bef += bef
         self.total += num_obs
-        self.data_range = torch.maximum(self.data_range, torch.max(target) - torch.min(target))
 
     def compute(self) -> Tensor:
         """Compute peak signal-to-noise ratio over state."""
@@ -120,7 +132,7 @@ def plot(
             >>> # Example plotting a single value
             >>> import torch
             >>> from torchmetrics.image import PeakSignalNoiseRatioWithBlockedEffect
-            >>> metric = PeakSignalNoiseRatioWithBlockedEffect()
+            >>> metric = PeakSignalNoiseRatioWithBlockedEffect(data_range=1.0)
             >>> metric.update(torch.rand(2, 1, 10, 10), torch.rand(2, 1, 10, 10))
             >>> fig_, ax_ = metric.plot()
 
@@ -130,7 +142,7 @@ def plot(
             >>> # Example plotting multiple values
             >>> import torch
             >>> from torchmetrics.image import PeakSignalNoiseRatioWithBlockedEffect
-            >>> metric = PeakSignalNoiseRatioWithBlockedEffect()
+            >>> metric = PeakSignalNoiseRatioWithBlockedEffect(data_range=1.0)
             >>> values = [ ]
             >>> for _ in range(10):
             ...     values.append(metric(torch.rand(2, 1, 10, 10), torch.rand(2, 1, 10, 10)))

tests/unittests/image/test_psnr.py

@@ -154,19 +154,10 @@ def test_reduction_for_dim_none(reduction):
     """Test that warnings are raised when then reduction parameter is combined with no dim provided arg."""
     match = f"The `reduction={reduction}` will not have any effect when `dim` is None."
     with pytest.warns(UserWarning, match=match):
-        PeakSignalNoiseRatio(reduction=reduction, dim=None)
+        PeakSignalNoiseRatio(data_range=10.0, reduction=reduction, dim=None)
 
     with pytest.warns(UserWarning, match=match):
-        peak_signal_noise_ratio(_inputs[0].preds, _inputs[0].target, reduction=reduction, dim=None)
-
-
-def test_missing_data_range():
-    """Check that error is raised if data range is not provided."""
-    with pytest.raises(ValueError, match="The `data_range` must be given when `dim` is not None."):
-        PeakSignalNoiseRatio(data_range=None, dim=0)
-
-    with pytest.raises(ValueError, match="The `data_range` must be given when `dim` is not None."):
-        peak_signal_noise_ratio(_inputs[0].preds, _inputs[0].target, data_range=None, dim=0)
+        peak_signal_noise_ratio(_inputs[0].preds, _inputs[0].target, data_range=10.0, reduction=reduction, dim=None)
 
 
 def test_psnr_uint_dtype():
@@ -177,6 +168,6 @@ def test_psnr_uint_dtype():
     """
     preds = torch.randint(0, 255, _input_size, dtype=torch.uint8)
     target = torch.randint(0, 255, _input_size, dtype=torch.uint8)
-    psnr = peak_signal_noise_ratio(preds, target)
-    prnr2 = peak_signal_noise_ratio(preds.float(), target.float())
+    psnr = peak_signal_noise_ratio(preds, target, data_range=255.0)
+    prnr2 = peak_signal_noise_ratio(preds.float(), target.float(), data_range=255.0)
     assert torch.allclose(psnr, prnr2)