fix: Handle mixed single/multi-channel y and y_pred in DiceHelper

Convert y_pred and y to boolean independently based on each tensor's
own channel count, fixing incorrect Dice values when formats differ
(e.g. single-channel class indices paired with multi-channel one-hot).

Add test cases covering both mixed-format combinations.

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>

Author: aymuos15

monai/metrics/meandice.py

@@ -307,23 +307,21 @@ def __call__(self, y_pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor | tupl
         batch_size = y_pred.shape[0]
         device = y_pred.device
 
-        # Convert to boolean for computation
+        # Convert y_pred to boolean (handle single-channel class indices vs multi-channel one-hot independently)
         if y_pred.shape[1] == 1 and n_pred_ch > 1:
-            # Single-channel class indices: convert to one-hot
             y_pred_bool = torch.zeros(batch_size, n_pred_ch, *y_pred.shape[2:], dtype=torch.bool, device=device)
-            y_bool = torch.zeros(batch_size, n_pred_ch, *y.shape[2:], dtype=torch.bool, device=device)
-
             for c in range(n_pred_ch):
                 y_pred_bool[:, c] = (y_pred[:, 0] == c)
-                y_bool[:, c] = (y[:, 0] == c)
         else:
-            # One-hot format: cast to bool
             y_pred_bool = y_pred.bool()
-            if y.shape[1] == 1 and y_pred.shape[1] > 1:
-                # Expand y to match y_pred channels
-                y_bool = (y == 1).expand(batch_size, n_pred_ch, *y.shape[2:])
-            else:
-                y_bool = y.bool()
+
+        # Convert y to boolean (independent of y_pred format)
+        if y.shape[1] == 1 and n_pred_ch > 1:
+            y_bool = torch.zeros(batch_size, n_pred_ch, *y.shape[2:], dtype=torch.bool, device=device)
+            for c in range(n_pred_ch):
+                y_bool[:, c] = (y[:, 0] == c)
+        else:
+            y_bool = y.bool()
 
         # Flatten spatial dimensions for vectorized computation: (batch, channels, -1)
         y_pred_flat = y_pred_bool.reshape(batch_size, n_pred_ch, -1).float()

tests/metrics/test_compute_meandice.py

@@ -251,15 +251,71 @@
 ]
 
 
+# single-channel y (class indices) with multi-channel y_pred (one-hot)
+TEST_CASE_MIXED_1 = [
+    {
+        "y_pred": torch.tensor(
+            [[[[0.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [1.0, 0.0]]]]
+        ),  # (1, 3, 2, 2) one-hot
+        "y": torch.tensor([[[[0.0, 1.0], [2.0, 1.0]]]]),  # (1, 1, 2, 2) class indices
+        "include_background": True,
+    },
+    # class 0: y_gt=[[1,0],[0,0]], y_pred=[[0,1],[0,0]] -> dice=0.0
+    # class 1: y_gt=[[0,1],[0,1]], y_pred=[[0,0],[0,1]] -> dice=2/3
+    # class 2: y_gt=[[0,0],[1,0]], y_pred=[[1,0],[1,0]] -> dice=2/3
+    [[0.0000, 0.6667, 0.6667]],
+]
+
+# single-channel y_pred (argmaxed, with num_classes) with multi-channel y (one-hot)
+TEST_CASE_MIXED_2 = [
+    {
+        "y_pred": torch.tensor([[[[2.0, 2.0], [2.0, 2.0]]]]),  # (1, 1, 2, 2) all class 2
+        "y": torch.tensor(
+            [[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]
+        ),  # (1, 3, 2, 2) one-hot, all background
+        "include_background": True,
+        "num_classes": 3,
+    },
+    # class 0: y_gt=[1,1,1,1](4), y_pred=[0,0,0,0](0) -> dice=0.0
+    # class 1: y_gt=[0,0,0,0](0), y_pred=[0,0,0,0](0) -> dice=nan (ignore_empty default)
+    # class 2: y_gt=[0,0,0,0](0), y_pred=[1,1,1,1](4) -> dice=nan (ignore_empty default)
+    [[False, True, True]],  # False=not-nan, True=nan
+]
+
+# single-channel y (class indices) with multi-channel y_pred, exclude background
+TEST_CASE_MIXED_3 = [
+    {
+        "y_pred": torch.tensor(
+            [
+                [[[0.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]]],
+                [[[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 0.0]]],
+            ]
+        ),  # (2, 3, 2, 2) one-hot
+        "y": torch.tensor(
+            [
+                [[[0.0, 0.0], [0.0, 1.0]]],
+                [[[0.0, 0.0], [0.0, 1.0]]],
+            ]
+        ),  # (2, 1, 2, 2) class indices
+        "include_background": False,
+    },
+    # batch 0: class 1 y_gt=[[0,0],[0,1]], y_pred=[[0,0],[1,1]] -> dice=2/3
+    #          class 2 y_gt=[[0,0],[0,0]], y_pred=[[1,0],[0,0]] -> dice=nan
+    # batch 1: class 1 y_gt=[[0,0],[0,1]], y_pred=[[1,0],[0,0]] -> dice=0.0
+    #          class 2 y_gt=[[0,0],[0,0]], y_pred=[[0,1],[1,0]] -> dice=nan
+    [[False, True], [False, True]],  # nan pattern
+]
+
+
 class TestComputeMeanDice(unittest.TestCase):
 
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_9, TEST_CASE_11, TEST_CASE_12])
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_9, TEST_CASE_11, TEST_CASE_12, TEST_CASE_MIXED_1])
     def test_value(self, input_data, expected_value):
         result = compute_dice(**input_data)
         np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4)
         np.testing.assert_equal(result.device, input_data["y_pred"].device)
 
-    @parameterized.expand([TEST_CASE_3])
+    @parameterized.expand([TEST_CASE_3, TEST_CASE_MIXED_2, TEST_CASE_MIXED_3])
     def test_nans(self, input_data, expected_value):
         result = compute_dice(**input_data)
         self.assertTrue(np.allclose(np.isnan(result.cpu().numpy()), expected_value))