Fix MS-SSIM size validation across scale counts

monai/metrics/regression.py

@@ -578,7 +578,8 @@ def compute_ms_ssim(
         k2: stability constant used in the contrast denominator
         weights: parameters for image similarity and contrast sensitivity at different resolution scores.
     Raises:
-        ValueError: when `y_pred` is not a 2D or 3D image.
+        ValueError: when `y_pred` is not a 2D or 3D image, or its spatial dimensions are too small
+            for the kernel size and number of scales.
     """
     dims = y_pred.ndimension()
     if spatial_dims == 2 and dims != 4:
@@ -598,15 +599,15 @@ def compute_ms_ssim(
 
     if not isinstance(kernel_sigma, Sequence):
         kernel_sigma = ensure_tuple_rep(kernel_sigma, spatial_dims)
-    # check if image have enough size for the number of downsamplings and the size of the kernel
-    weights_div = max(1, (len(weights) - 1)) ** 2
+    # Check that the last scale can accommodate the kernel after stride-2 downsampling.
+    weights_div = 2 ** (len(weights) - 1)
     y_pred_spatial_dims = y_pred.shape[2:]
     for i in range(len(y_pred_spatial_dims)):
-        if y_pred_spatial_dims[i] // weights_div <= kernel_size[i] - 1:
+        if y_pred_spatial_dims[i] // weights_div < kernel_size[i]:
             raise ValueError(
                 f"For a given number of `weights` parameters {len(weights)} and kernel size "
-                f"{kernel_size[i]}, the image height must be larger than "
-                f"{(kernel_size[i] - 1) * weights_div}."
+                f"{kernel_size[i]}, spatial dimension {i} of the image must be at least "
+                f"{kernel_size[i] * weights_div}."
             )
 
     weights_tensor = torch.tensor(weights, device=y_pred.device, dtype=torch.float)

tests/metrics/test_compute_multiscalessim_metric.py

@@ -14,6 +14,7 @@
 import unittest
 
 import torch
+from parameterized import parameterized
 
 from monai.metrics import MultiScaleSSIMMetric
 from monai.utils import set_determinism
@@ -72,6 +73,47 @@ def test_input_ill_input_shape3d(self):
         with self.assertRaises(ValueError):
             metric(torch.randn(1, 1, 64, 64, 64), torch.randn(1, 1, 64, 64, 64))
 
+    @parameterized.expand([(spatial_dims, scales) for spatial_dims in (2, 3) for scales in range(1, 7)])
+    def test_minimum_input_size(self, spatial_dims, scales):
+        kernel_size = 3
+        minimum_size = kernel_size * 2 ** (scales - 1)
+        image = torch.ones(1, 1, *([minimum_size] * spatial_dims))
+        metric = MultiScaleSSIMMetric(
+            spatial_dims=spatial_dims, kernel_size=kernel_size, weights=[1.0 / scales] * scales
+        )
+
+        self.assertAlmostEqual(metric(image, image).item(), 1.0)
+
+    @parameterized.expand([(spatial_dims, scales) for spatial_dims in (2, 3) for scales in range(1, 7)])
+    def test_undersized_spatial_dimension(self, spatial_dims, scales):
+        kernel_size = (3, 5) if spatial_dims == 2 else (3, 5, 3)
+        minimum_shape = [size * 2 ** (scales - 1) for size in kernel_size]
+        metric = MultiScaleSSIMMetric(
+            spatial_dims=spatial_dims, kernel_size=kernel_size, weights=[1.0 / scales] * scales
+        )
+
+        for axis in range(spatial_dims):
+            with self.subTest(axis=axis):
+                shape = minimum_shape.copy()
+                shape[axis] -= 1
+                image = torch.ones(1, 1, *shape)
+                message = f"spatial dimension {axis} of the image must be at least {minimum_shape[axis]}"
+                with self.assertRaisesRegex(ValueError, message):
+                    metric(image, image)
+
+    def test_four_scales_default_kernel(self):
+        image = torch.ones(1, 1, 88, 88)
+        metric = MultiScaleSSIMMetric(spatial_dims=2, weights=[0.25] * 4)
+
+        self.assertAlmostEqual(metric(image, image).item(), 1.0)
+
+    def test_two_scales_too_small(self):
+        image = torch.ones(1, 1, 16, 16)
+        metric = MultiScaleSSIMMetric(spatial_dims=2, weights=[0.5, 0.5])
+
+        with self.assertRaisesRegex(ValueError, "spatial dimension 0 of the image must be at least 22"):
+            metric(image, image)
+
     def test_small_inputs(self):
         metric = MultiScaleSSIMMetric(spatial_dims=2)