Repository: Lightning-AI/torchmetrics

Issue:
Issues with Spectral Distortion Index
## 🐛 Bug

I'm having a few issues with the implementation of Spectral Distortion Index.

1) Since SDI is a reference free metric for super-resolution, the inputs are usually not going to be the same size, but the function throws an error if the resolutions are different. I believe only the batch and channel sizes should be validated.
2) For some inputs, I'm getting `nan` value outputs. In a set of 49 hyperspectral images I'm working with, I have 14 that result in `nan`. I believe this happens whenever the denominator `lower` is 0 here:
https://github.com/Lightning-AI/metrics/blob/780e9afa7316ab4439b888304eb5c14fdbaacbef/src/torchmetrics/functional/image/uqi.py#L117
This is a known issue with UQI, which is the reason it was adapted into SSIM by adding a constant to the denominator. Would it make sense to add an option to use SSIM instead of UQI with SDI?
3) As a result of the nested for loops, the implementation is very slow to run for images with many channels - I'm not sure that they can easily be vectorised, unfortunately, so maybe there's not much that can be done about this. For reference, here are execution times for the aforementioned test set:

> PSNR: 32.4233 (total execution time 2.21s)
> SSIM: 0.8302 (total execution time 27.60s)
> MS-SSIM: 0.9573 (total execution time 35.06s)
> SAM: 0.0819 (total execution time 1.37s)
> SDI: nan (total execution time 16694.23s)

### To Reproduce

Here's a MWE with an image that results in a `nan` value:

[sdi_test.tar.gz](https://github.com/Lightning-AI/metrics/files/10767083/sdi_test.tar.gz)

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): 0.12.0dev, built from source
- Python & PyTorch Version (e.g., 1.0): Python 3.10.6, PyTorch 1.13.1+cu117
- Any other relevant information such as OS (e.g., Linux): Linux

### Additional context

I have a fix for 1) here: https://github.com/rddunphy/metrics/commit/ffcea1570a0c4da0607589c9d12283cf9a79a05d

Let me know if I should open a PR, or have a go at fixing 2) first. Thanks!


Discussion:
Hi! thanks for your contribution!, great first issue!
Hi @rddunphy,
Please feel free to send the fixes for 1) and also feel free to include the proposed changes for 2). I do not see a problem with adding a small epsilon to the denominator to make sure we do not divide by 0. 

For 3), it is possible to get rid of the inner for-loop fairly easy by stacking the input:
```python
        stack1 = torch.cat([target[:, k : k + 1, :, :] for r in range(k, length)], dim=0)
        stack2 = torch.cat([target[:, r : r + 1, :, :] for r in range(k, length)], dim=0)
```
then running `universal_image_quality_index` on the stack with `reduction="none"` and afterwards do the averaging manually. This would most likely be faster but also require much more memory because we would need to store input multiple times in memory. So not sure what to do it in this case.

Resolving patch:
diff --git a/src/torchmetrics/functional/image/d_lambda.py b/src/torchmetrics/functional/image/d_lambda.py
--- a/src/torchmetrics/functional/image/d_lambda.py
+++ b/src/torchmetrics/functional/image/d_lambda.py
@@ -19,7 +19,6 @@
 from typing_extensions import Literal
 
 from torchmetrics.functional.image.uqi import universal_image_quality_index
-from torchmetrics.utilities.checks import _check_same_shape
 from torchmetrics.utilities.distributed import reduce
 
 
@@ -34,11 +33,15 @@ def _spectral_distortion_index_update(preds: Tensor, target: Tensor) -> Tuple[Te
         raise TypeError(
             f"Expected `ms` and `fused` to have the same data type. Got ms: {preds.dtype} and fused: {target.dtype}."
         )
-    _check_same_shape(preds, target)
     if len(preds.shape) != 4:
         raise ValueError(
             f"Expected `preds` and `target` to have BxCxHxW shape. Got preds: {preds.shape} and target: {target.shape}."
         )
+    if preds.shape[:2] != target.shape[:2]:
+        raise ValueError(
+            "Expected `preds` and `target` to have same batch and channel sizes."
+            f"Got preds: {preds.shape} and target: {target.shape}."
+        )
     return preds, target
 
 
@@ -69,13 +72,29 @@ def _spectral_distortion_index_compute(
         tensor(0.0234)
     """
     length = preds.shape[1]
-    m1 = torch.zeros((length, length))
-    m2 = torch.zeros((length, length))
+
+    m1 = torch.zeros((length, length), device=preds.device)
+    m2 = torch.zeros((length, length), device=preds.device)
 
     for k in range(length):
-        for r in range(k, length):
-            m1[k, r] = m1[r, k] = universal_image_quality_index(target[:, k : k + 1, :, :], target[:, r : r + 1, :, :])
-            m2[k, r] = m2[r, k] = universal_image_quality_index(preds[:, k : k + 1, :, :], preds[:, r : r + 1, :, :])
+        num = length - (k + 1)
+        if num == 0:
+            continue
+        stack1 = target[:, k : k + 1, :, :].repeat(num, 1, 1, 1)
+        stack2 = torch.cat([target[:, r : r + 1, :, :] for r in range(k + 1, length)], dim=0)
+        score = [
+            s.mean() for s in universal_image_quality_index(stack1, stack2, reduction="none").split(preds.shape[0])
+        ]
+        m1[k, k + 1 :] = torch.stack(score, 0)
+
+        stack1 = preds[:, k : k + 1, :, :].repeat(num, 1, 1, 1)
+        stack2 = torch.cat([preds[:, r : r + 1, :, :] for r in range(k + 1, length)], dim=0)
+        score = [
+            s.mean() for s in universal_image_quality_index(stack1, stack2, reduction="none").split(preds.shape[0])
+        ]
+        m2[k, k + 1 :] = torch.stack(score, 0)
+    m1 = m1 + m1.T
+    m2 = m2 + m2.T
 
     diff = torch.pow(torch.abs(m1 - m2), p)
     # Special case: when number of channels (L) is 1, there will be only one element in M1 and M2. Hence no need to sum.
diff --git a/src/torchmetrics/functional/image/uqi.py b/src/torchmetrics/functional/image/uqi.py
--- a/src/torchmetrics/functional/image/uqi.py
+++ b/src/torchmetrics/functional/image/uqi.py
@@ -112,7 +112,7 @@ def _uqi_compute(
     sigma_pred_target = output_list[4] - mu_pred_target
 
     upper = 2 * sigma_pred_target
-    lower = sigma_pred_sq + sigma_target_sq
+    lower = sigma_pred_sq + sigma_target_sq + torch.finfo(sigma_pred_sq.dtype).eps
 
     uqi_idx = ((2 * mu_pred_target) * upper) / ((mu_pred_sq + mu_target_sq) * lower)
     uqi_idx = uqi_idx[..., pad_h:-pad_h, pad_w:-pad_w]


Test patch:
diff --git a/tests/unittests/image/test_d_lambda.py b/tests/unittests/image/test_d_lambda.py
--- a/tests/unittests/image/test_d_lambda.py
+++ b/tests/unittests/image/test_d_lambda.py
@@ -118,12 +118,6 @@ def test_d_lambda_functional(self, preds, target, p):
             metric_args={"p": p},
         )
 
-    # SpectralDistortionIndex half + cpu does not work due to missing support in torch.log
-    @pytest.mark.xfail(reason="Spectral Distortion Index metric does not support cpu + half precision")
-    def test_d_lambda_half_cpu(self, preds, target, p):
-        """Test dtype support of the metric on CPU."""
-        self.run_precision_test_cpu(preds, target, SpectralDistortionIndex, spectral_distortion_index, {"p": p})
-
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires cuda")
     def test_d_lambda_half_gpu(self, preds, target, p):
         """Test dtype support of the metric on GPU."""
@@ -152,3 +146,11 @@ def test_d_lambda_invalid_type():
     target_t = torch.rand((1, 1, 16, 16), dtype=torch.float64)
     with pytest.raises(TypeError, match="Expected `ms` and `fused` to have the same data type.*"):
         spectral_distortion_index(preds_t, target_t, p=1)
+
+
+def test_d_lambda_different_sizes():
+    """Since d lambda is reference free, it can accept different number of targets and preds."""
+    preds = torch.rand(1, 1, 32, 32)
+    target = torch.rand(1, 1, 16, 16)
+    out = spectral_distortion_index(preds, target, p=1)
+    assert isinstance(out, torch.Tensor)