fix(metrics): exclude ignored voxels from confusion matrix scores

monai/metrics/confusion_matrix.py

@@ -69,6 +69,7 @@ class ConfusionMatrixMetric(CumulativeIterationMetric):
             Voxels with this label are excluded from the score, which is useful for padding, unlabeled regions,
             or boundary artifacts. For federated or aggregated settings, ensure all clients use the same
             ignore_index to keep score values comparable.
+            For a single-channel binary target, ignoring channel 0 excludes locations where the target is 1.
 
     """
 
@@ -159,11 +160,19 @@ def get_confusion_matrix(
             The values should be binarized.
         y: ground truth to compute the metric. It must be one-hot format and first dim is batch.
             The values should be binarized.
+            A single channel is a binary membership indicator, not a map of class IDs.
         include_background: whether to include metric computation on the first channel of
             the predicted output. Defaults to True.
         ignore_index: single integer class index (or sentinel value) to ignore from the metric computation.
-            If ignore_index refers to a valid class channel, that channel is excluded. If it is outside the
-            class range, ignored regions are inferred from spatial locations where all label channels are zero.
+            For a valid class index, voxels belonging to that channel are excluded from all classes' counts.
+            The existing channel-exclusion behavior is retained for multi-channel targets. With a single-channel
+            binary target, ignore_index=0 excludes locations where the target is 1 without zeroing the channel.
+            Counts at valid single-channel locations do not depend on include_background. Out-of-range values use all-zero
+            locations for multi-channel targets and direct value comparison for single-channel targets.
+
+    Returns:
+        Tensor of shape [B, C, 4], containing true positive, false positive, true negative, and false negative
+        counts, in that order, for each sample and evaluated channel.
 
     Raises:
         ValueError: when `y_pred` and `y` have different shapes.
@@ -179,17 +188,23 @@ class range, ignored regions are inferred from spatial locations where all label
     # get confusion matrix related metric
     batch_size, n_class = y_pred.shape[:2]
 
-    # Create spatial mask if ignore_index is provided
-    if ignore_index is not None and 0 <= ignore_index < (y_pred.shape[1] + (0 if include_background else 1)):
+    # Single-channel targets need spatial exclusion only, preserving errors at valid locations.
+    if (
+        ignore_index is not None
+        and original_y.shape[1] > 1
+        and 0 <= ignore_index < (y_pred.shape[1] + (0 if include_background else 1))
+    ):
         ignore_channel = ignore_index if include_background else ignore_index - 1
         if 0 <= ignore_channel < y_pred.shape[1]:
             y_pred = y_pred.clone()
             y = y.clone()
             y_pred[:, ignore_channel] = 0
             y[:, ignore_channel] = 0
-        mask = None
+    if ignore_index is not None and 0 <= ignore_index < original_y.shape[1]:
+        # Every channel represents binary membership; comparison also supports boolean targets.
+        mask = (original_y[:, ignore_index : ignore_index + 1] == 0).float()
     else:
-        mask = create_ignore_mask(original_y if ignore_index is not None else y, ignore_index)
+        mask = create_ignore_mask(original_y, ignore_index)
 
     # convert to [BNS], where S is the number of pixels for one sample.
     y_pred = y_pred.reshape(batch_size, n_class, -1)

tests/metrics/test_ignore_index_metrics.py

@@ -26,6 +26,7 @@
     SurfaceDistanceMetric,
     compute_dice,
     compute_iou,
+    get_confusion_matrix,
 )
 from monai.utils import optional_import
 
@@ -158,6 +159,173 @@ def test_ignored_voxels_excluded_from_other_classes(self):
         self.assertEqual(iou[0, 0].item(), 1.0)
         torch.testing.assert_close(iou, dice, equal_nan=True)
 
+    @parameterized.expand([(background, sample) for background in (True, False) for sample in (True, False)])
+    def test_confusion_matrix_excludes_ignored_voxels(self, include_background, compute_sample):
+        """Exclude ignored-class voxels from counts and precision in both aggregation modes.
+
+        Args:
+            include_background: whether to include the background channel in the output.
+            compute_sample: whether to compute precision per sample before aggregation.
+        """
+        y = torch.tensor([[[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]]])
+        # The only error predicts class 2 at a voxel belonging to ignored class 1.
+        y_pred = torch.tensor([[[1.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 1.0, 0.0]]])
+        metric = ConfusionMatrixMetric(
+            include_background=include_background,
+            metric_name="precision",
+            compute_sample=compute_sample,
+            reduction="none",
+            ignore_index=1,
+        )
+        metric(y_pred=y_pred, y=y)
+        class_index = 2 if include_background else 1
+        self.assertEqual(metric.aggregate()[0][0, class_index].item(), 1.0)
+        counts = get_confusion_matrix(y_pred, y, include_background=include_background, ignore_index=1)
+        self.assertEqual(counts[0, class_index].tolist(), [1.0, 0.0, 2.0, 0.0])
+
+    @parameterized.expand([(True,), (False,)])
+    def test_confusion_matrix_excludes_ignored_background_from_negatives(self, include_background):
+        """Exclude ignored background voxels from every channel's negative counts.
+
+        Args:
+            include_background: whether to include the background channel in the output.
+        """
+        y = torch.tensor([[[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]]])
+        y_pred = y.clone()
+        y_pred[0, 0, 0] = 0.0
+        y_pred[0, 1, 0] = 1.0
+        counts = get_confusion_matrix(y_pred, y, include_background=include_background, ignore_index=0)
+        class_index = 2 if include_background else 1
+        self.assertEqual(counts[0, class_index].tolist(), [1.0, 0.0, 1.0, 0.0])
+
+    def test_confusion_matrix_keeps_errors_on_valid_voxels(self):
+        """Preserve false positives and false negatives at valid voxels."""
+        y = torch.tensor([[[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]]])
+        y_pred = torch.tensor([[[0.0, 0.0, 0.0, 1.0], [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 1.0, 0.0]]])
+        counts = get_confusion_matrix(y_pred, y, include_background=False, ignore_index=0)
+        self.assertEqual(counts[0, 0].tolist(), [0.0, 0.0, 1.0, 1.0])
+        self.assertEqual(counts[0, 1].tolist(), [1.0, 1.0, 0.0, 0.0])
+
+    @parameterized.expand([(True,), (False,)])
+    def test_confusion_matrix_sentinel_mask(self, include_background):
+        """Preserve sentinel masking with and without the background channel.
+
+        Args:
+            include_background: whether to include the background channel in the output.
+        """
+        y = torch.tensor([[[1.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]]])
+        y_pred = torch.tensor([[[1.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 1.0, 0.0]]])
+        counts = get_confusion_matrix(y_pred, y, include_background=include_background, ignore_index=255)
+        class_index = 2 if include_background else 1
+        self.assertEqual(counts[0, class_index].tolist(), [1.0, 0.0, 2.0, 0.0])
+
+    @parameterized.expand(
+        [
+            (background, ignore_index, dtype)
+            for background in (True, False)
+            for ignore_index in (0, 1, 2)
+            for dtype in (torch.float32, torch.int64, torch.bool)
+        ]
+    )
+    def test_confusion_matrix_multichannel_boolean(self, include_background, ignore_index, prediction_dtype):
+        """Exclude ignored voxels from boolean one-hot targets without mutating inputs.
+
+        Args:
+            include_background: whether to retain the background channel in the returned counts.
+            ignore_index: original class channel whose spatial locations must be excluded.
+            prediction_dtype: data type of the binarized prediction tensor.
+        """
+        y = torch.tensor([[[1, 0, 0, 1], [0, 1, 0, 0], [0, 0, 1, 0]]], dtype=torch.bool)
+        y_pred = torch.tensor([[[1, 1, 0, 1], [0, 0, 0, 0], [0, 0, 1, 0]]], dtype=prediction_dtype)
+        original_y, original_pred = y.clone(), y_pred.clone()
+        expected = torch.tensor(
+            [
+                [[0, 0, 2, 0], [0, 0, 1, 1], [1, 0, 1, 0]],
+                [[2, 0, 1, 0], [0, 0, 3, 0], [1, 0, 2, 0]],
+                [[2, 1, 0, 0], [0, 0, 2, 1], [0, 0, 3, 0]],
+            ],
+            dtype=torch.float32,
+        )[ignore_index].unsqueeze(0)
+        if not include_background:
+            expected = expected[:, 1:]
+
+        counts = get_confusion_matrix(y_pred, y, include_background=include_background, ignore_index=ignore_index)
+
+        torch.testing.assert_close(counts, expected)
+        torch.testing.assert_close(y, original_y)
+        torch.testing.assert_close(y_pred, original_pred)
+
+    @parameterized.expand(
+        [(background, dtype) for background in (True, False) for dtype in (torch.float32, torch.int64, torch.bool)]
+    )
+    def test_confusion_matrix_single_channel_ignore_class(self, include_background, dtype):
+        """Interpret a single channel as binary membership when ignoring class zero.
+
+        Args:
+            include_background: whether to request inclusion of the background channel.
+            dtype: data type of the binary target and prediction tensors.
+        """
+        y = torch.tensor([[[1, 1, 0]]], dtype=dtype)
+        y_pred = torch.ones_like(y)
+        counts = get_confusion_matrix(y_pred, y, include_background=include_background, ignore_index=0)
+        self.assertEqual(counts[0, 0].tolist(), [0.0, 1.0, 0.0, 0.0])
+
+    @parameterized.expand([(True,), (False,)])
+    def test_confusion_matrix_single_channel_specificity(self, include_background):
+        """Retain a false positive at the only valid single-channel voxel.
+
+        Args:
+            include_background: whether to request inclusion of the background channel.
+        """
+        y = torch.tensor([[[1.0, 1.0, 0.0]]])
+        metric = ConfusionMatrixMetric(
+            include_background=include_background, metric_name="specificity", reduction="none", ignore_index=0
+        )
+        metric(y_pred=torch.ones_like(y), y=y)
+        self.assertEqual(metric.aggregate()[0].item(), 0.0)
+
+    @parameterized.expand(
+        [
+            (True, [1.0, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0]),
+            (False, [1.0, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0]),
+            (True, [0.0, 0.0, 0.0], [0.0, 3.0, 0.0, 0.0]),
+            (False, [0.0, 0.0, 0.0], [0.0, 3.0, 0.0, 0.0]),
+        ]
+    )
+    def test_confusion_matrix_single_channel_mask_extremes(self, include_background, labels, expected):
+        """Handle all-ignored and entirely valid single-channel spatial regions.
+
+        Args:
+            include_background: whether to request inclusion of the background channel.
+            labels: binary target values for the spatial locations.
+            expected: expected counts in true positive, false positive, true negative, false negative order.
+        """
+        y = torch.tensor([[labels]])
+        counts = get_confusion_matrix(torch.ones_like(y), y, include_background=include_background, ignore_index=0)
+        self.assertEqual(counts[0, 0].tolist(), expected)
+
+    @parameterized.expand([(True,), (False,)])
+    def test_confusion_matrix_single_channel_sentinel(self, include_background):
+        """Keep the existing out-of-range comparison for single-channel targets.
+
+        Args:
+            include_background: whether to request inclusion of the background channel.
+        """
+        y = torch.tensor([[[1.0, 1.0, 0.0]]])
+        counts = get_confusion_matrix(torch.ones_like(y), y, include_background=include_background, ignore_index=255)
+        self.assertEqual(counts[0, 0].tolist(), [2.0, 1.0, 0.0, 0.0])
+
+    @parameterized.expand([(True,), (False,)])
+    def test_confusion_matrix_single_channel_without_ignore(self, include_background):
+        """Leave the unmasked single-channel counts unchanged.
+
+        Args:
+            include_background: whether to request inclusion of the background channel.
+        """
+        y = torch.tensor([[[1.0, 1.0, 0.0]]])
+        counts = get_confusion_matrix(torch.ones_like(y), y, include_background=include_background)
+        self.assertEqual(counts[0, 0].tolist(), [2.0, 1.0, 0.0, 0.0])
+
     def test_ignored_voxels_excluded_with_include_background_false(self):
         """The ignore_index mask must line up with the ignore_background channel strip."""
         # 4 one-hot classes: 0=background, 1, 2=ignored, 3