Fix overflow for PSNR metric when used with uint8 input (#2788)

CHANGELOG.md

@@ -56,6 +56,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed corner case in `Iou` metric for single empty prediction tensors ([#2780](https://github.com/Lightning-AI/torchmetrics/pull/2780))
 
 
+- Fixed `PSNR` calculation for integer type input images ([#2788](https://github.com/Lightning-AI/torchmetrics/pull/2788))
+
+
 ## [1.4.3] - 2024-10-10
 
 ### Fixed

src/torchmetrics/functional/image/psnr.py

@@ -69,6 +69,11 @@ def _psnr_update(
             Default is None meaning scores will be reduced across all dimensions.
 
     """
+    if not preds.is_floating_point():
+        preds = preds.to(torch.float32)
+    if not target.is_floating_point():
+        target = target.to(torch.float32)
+
     if dim is None:
         sum_squared_error = torch.sum(torch.pow(preds - target, 2))
         num_obs = tensor(target.numel(), device=target.device)

tests/unittests/image/test_psnr.py

@@ -167,3 +167,16 @@ def test_missing_data_range():
 
     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)
+
+
+def test_psnr_uint_dtype():
+    """Check that automatic casting to float is done for uint dtype.
+
+    See issue: https://github.com/Lightning-AI/torchmetrics/issues/2787
+
+    """
+    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())
+    assert torch.allclose(psnr, prnr2)