Fix reduction for full image in SSIM

CHANGELOG.md

@@ -43,7 +43,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
--
+- Fixed a bug in `ssim` when `return_full_image=True` where the score was still reduced ([#1204](https://github.com/Lightning-AI/metrics/pull/1204))
 
 
 -

src/torchmetrics/functional/image/ssim.py

@@ -69,7 +69,7 @@ def _ssim_compute(
             Ignored if a uniform kernel is used
         kernel_size: the size of the uniform kernel, anisotropic kernels are possible.
             Ignored if a Gaussian kernel is used
-        reduction: a method to reduce metric score over labels.
+        reduction: a method to reduce metric score over individual batch scores
 
             - ``'elementwise_mean'``: takes the mean
             - ``'sum'``: takes the sum
@@ -117,6 +117,9 @@ def _ssim_compute(
             f"Expected `kernel_size` dimension to be 2 or 3. `kernel_size` dimensionality: {len(kernel_size)}"
         )
 
+    if return_full_image and return_contrast_sensitivity:
+        raise ValueError("Arguments `return_full_image` and `return_contrast_sensitivity` are mutually exclusive.")
+
     if any(x % 2 == 0 or x <= 0 for x in kernel_size):
         raise ValueError(f"Expected `kernel_size` to have odd positive number. Got {kernel_size}.")
 
@@ -189,9 +192,7 @@ def _ssim_compute(
         )
 
     elif return_full_image:
-        return reduce(ssim_idx.reshape(ssim_idx.shape[0], -1).mean(-1), reduction), reduce(
-            ssim_idx_full_image, reduction
-        )
+        return reduce(ssim_idx.reshape(ssim_idx.shape[0], -1).mean(-1), reduction), ssim_idx_full_image
 
     return reduce(ssim_idx.reshape(ssim_idx.shape[0], -1).mean(-1), reduction)
 

src/torchmetrics/image/ssim.py

@@ -33,7 +33,7 @@ class StructuralSimilarityIndexMeasure(Metric):
             Ignored if a uniform kernel is used
         kernel_size: the size of the uniform kernel, anisotropic kernels are possible.
             Ignored if a Gaussian kernel is used
-        reduction: a method to reduce metric score over labels.
+        reduction: a method to reduce metric score over individual batch scores
 
             - ``'elementwise_mean'``: takes the mean
             - ``'sum'``: takes the sum

tests/unittests/image/test_ssim.py

@@ -18,6 +18,7 @@
 import pytest
 import torch
 from skimage.metrics import structural_similarity
+from torch import Tensor
 
 from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image import StructuralSimilarityIndexMeasure
@@ -52,7 +53,16 @@
     )
 
 
-def _sk_ssim(preds, target, data_range, sigma, kernel_size=None, return_ssim_image=False, gaussian_weights=True):
+def _sk_ssim(
+    preds,
+    target,
+    data_range,
+    sigma,
+    kernel_size=None,
+    return_ssim_image=False,
+    gaussian_weights=True,
+    reduction_arg="elementwise_mean",
+):
     if len(preds.shape) == 4:
         c, h, w = preds.shape[-3:]
         sk_preds = preds.view(-1, c, h, w).permute(0, 2, 3, 1).numpy()
@@ -77,7 +87,7 @@ def _sk_ssim(preds, target, data_range, sigma, kernel_size=None, return_ssim_ima
                 full=return_ssim_image,
             )
             results[i] = torch.from_numpy(np.asarray(res)).type(preds.dtype)
-        return results
+        return results if reduction_arg != "sum" else results.sum()
     else:
         fullimages = torch.zeros(target.shape, dtype=target.dtype)
         for i in range(sk_preds.shape[0]):
@@ -142,13 +152,14 @@ def test_ssim_without_gaussian_kernel(self, preds, target, sigma, ddp, dist_sync
             dist_sync_on_step=dist_sync_on_step,
         )
 
-    def test_ssim_functional(self, preds, target, sigma):
+    @pytest.mark.parametrize("reduction_arg", ["sum", "elementwise_mean", None])
+    def test_ssim_functional(self, preds, target, sigma, reduction_arg):
         self.run_functional_metric_test(
             preds,
             target,
             structural_similarity_index_measure,
-            partial(_sk_ssim, data_range=1.0, sigma=sigma, kernel_size=None),
-            metric_args={"data_range": 1.0, "sigma": sigma},
+            partial(_sk_ssim, data_range=1.0, sigma=sigma, kernel_size=None, reduction_arg=reduction_arg),
+            metric_args={"data_range": 1.0, "sigma": sigma, "reduction": reduction_arg},
         )
 
     # SSIM half + cpu does not work due to missing support in torch.log
@@ -241,3 +252,19 @@ def test_ssim_unequal_kernel_size():
         structural_similarity_index_measure(preds, target, gaussian_kernel=False, kernel_size=(5, 3)),
         torch.tensor(0.05131844),
     )
+
+
+@pytest.mark.parametrize(
+    "preds, target",
+    [(i.preds, i.target) for i in _inputs],
+)
+def test_full_image_output(preds, target):
+    out = structural_similarity_index_measure(preds[0], target[0])
+    assert isinstance(out, Tensor)
+    assert out.numel() == 1
+
+    out = structural_similarity_index_measure(preds[0], target[0], return_full_image=True)
+    assert isinstance(out, tuple)
+    assert len(out) == 2
+    assert out[0].numel() == 1
+    assert out[1].shape == preds[0].shape