diff --git a/monai/metrics/meandice.py b/monai/metrics/meandice.py index 5320f0c1b4..089e292466 100644 --- a/monai/metrics/meandice.py +++ b/monai/metrics/meandice.py @@ -331,13 +331,17 @@ def __init__( self.per_component = per_component self.ignore_index = ignore_index - def compute_cc_dice(self, y_pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + def compute_cc_dice(self, y_pred: torch.Tensor, y: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: """ Compute per-component Dice for a single batch item. Args: y_pred (torch.Tensor): Predictions with shape (1, 2, D, H, W) or (1, 2, H, W). y (torch.Tensor): Ground truth with shape (1, 2, D, H, W) or (1, 2, H, W). + mask (torch.Tensor, optional): spatial mask (1, D, H, W) or (1, H, W) from + ``create_ignore_mask``, 1=valid, 0=ignore. Ignored voxels are forced to + agree as background in both `y` and `y_pred`, so they never contribute to + a component's TP/FP/FN, matching how `ignore_index` is handled elsewhere. Returns: torch.Tensor: Mean Dice over connected components. @@ -348,6 +352,10 @@ def compute_cc_dice(self, y_pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor else: y_pred_idx = y_pred y_idx = y + if mask is not None: + ignore = mask == 0 + y_idx = y_idx.masked_fill(ignore, 0) + y_pred_idx = y_pred_idx.masked_fill(ignore, 0) if y_idx[0].sum() == 0: if self.ignore_empty: data = torch.tensor(float("nan"), device=y_idx.device) @@ -445,7 +453,10 @@ def __call__(self, y_pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor | tupl data = [] for b in range(y_pred.shape[0]): if self.per_component: - data.append(self.compute_cc_dice(y_pred=y_pred[b].unsqueeze(0), y=y[b].unsqueeze(0)).reshape(-1)) + b_mask = mask[b, 0] if mask is not None else None + data.append( + self.compute_cc_dice(y_pred=y_pred[b].unsqueeze(0), y=y[b].unsqueeze(0), mask=b_mask).reshape(-1) + ) continue c_list = [] for c in range(first_ch, n_pred_ch) if n_pred_ch > 1 else [1]: diff --git a/tests/metrics/test_compute_meandice.py b/tests/metrics/test_compute_meandice.py index 1a000b29b6..b13d3f2b4b 100644 --- a/tests/metrics/test_compute_meandice.py +++ b/tests/metrics/test_compute_meandice.py @@ -290,6 +290,19 @@ [[0.6400], [1.0], [1.0], [1.0], [1.0]], ] +# Testcase for per_component DiceMetric with ignore_index: a disagreement confined +# entirely to the ignored (sentinel, all-zero one-hot) region must not affect the score. +y = torch.zeros((1, 2, 8, 8)) +y_hat = torch.zeros((1, 2, 8, 8)) +y[0, 1, :, :4] = 1 +y_hat[0, 1, :, :4] = 1 +y_hat[0, 1, :, 4:] = 1 # disagreement, but this region is ignored below +y[0, 0] = 1 - y[0, 1] +y[0, 0, :, 4:] = 0 # sentinel: both channels 0 -> ignored by ignore_index=2 +y_hat[0, 0] = 1 - y_hat[0, 1] + +TEST_CASE_18 = [{"per_component": True, "ignore_empty": False, "ignore_index": 2}, {"y": y, "y_pred": y_hat}, [[1.0]]] + class TestComputeMeanDice(unittest.TestCase): @@ -342,7 +355,7 @@ def test_nans_class(self, params, input_data, expected_value): np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4) # CC DiceMetric tests - @parameterized.expand([TEST_CASE_16, TEST_CASE_17]) + @parameterized.expand([TEST_CASE_16, TEST_CASE_17, TEST_CASE_18]) @unittest.skipUnless(has_ndimage, "Requires scipy.ndimage.") def test_cc_dice_value_nogpu(self, params, input_data, expected_value): dice_metric = DiceMetric(**params)