refactor: centralize ignore_index masking into create_ignore_mask helper

Signed-off-by: Rusheel Sharma <rusheelhere@gmail.com>

Author: Rusheel86

monai/losses/dice.py

@@ -25,6 +25,7 @@
 from monai.losses.utils import compute_tp_fp_fn
 from monai.networks import one_hot
 from monai.utils import DiceCEReduction, LossReduction, Weight, look_up_option
+from monai.metrics.utils import create_ignore_mask
 
 
 class DiceLoss(_Loss):
@@ -166,16 +167,15 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         if self.other_act is not None:
             input = self.other_act(input)
 
-        mask: torch.Tensor | None = None
-        if self.ignore_index is not None:
-            mask = (target != self.ignore_index).float()
+        original_target = target if self.ignore_index is not None else None
 
         if self.to_onehot_y:
             if n_pred_ch == 1:
                 warnings.warn("single channel prediction, `to_onehot_y=True` ignored.")
             else:
                 target = one_hot(target, num_classes=n_pred_ch)
 
+        mask = create_ignore_mask(original_target if original_target is not None else target, self.ignore_index)
         if not self.include_background:
             if n_pred_ch == 1:
                 warnings.warn("single channel prediction, `include_background=False` ignored.")

monai/losses/focal_loss.py

@@ -20,6 +20,7 @@
 
 from monai.networks import one_hot
 from monai.utils import LossReduction
+from monai.metrics.utils import create_ignore_mask
 
 
 class FocalLoss(_Loss):
@@ -164,8 +165,8 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         if target.shape != input.shape:
             raise ValueError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
 
-        if self.ignore_index is not None:
-            mask = (target != self.ignore_index).float()
+        mask = create_ignore_mask(target, self.ignore_index)
+        if mask is not None:
             input = input * mask
             target = target * mask
 

monai/losses/tversky.py

@@ -20,6 +20,7 @@
 from monai.losses.utils import compute_tp_fp_fn
 from monai.networks import one_hot
 from monai.utils import LossReduction
+from monai.metrics.utils import create_ignore_mask
 
 
 class TverskyLoss(_Loss):
@@ -137,15 +138,10 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
 
         if self.ignore_index is not None:
             mask_src = original_target if self.to_onehot_y and n_pred_ch > 1 else target
-
-            if mask_src.shape[1] == 1:
-                mask = (mask_src != self.ignore_index).to(input.dtype)
-            else:
-                # Fallback for cases where target is already one-hot
-                mask = (1.0 - mask_src[:, self.ignore_index : self.ignore_index + 1]).to(input.dtype)
-
-            input = input * mask
-            target = target * mask
+            mask = create_ignore_mask(mask_src, self.ignore_index)
+            if mask is not None:
+                input = input * mask
+                target = target * mask
 
         if not self.include_background:
             if n_pred_ch == 1:

monai/losses/unified_focal_loss.py

@@ -18,6 +18,7 @@
 
 from monai.networks import one_hot
 from monai.utils import LossReduction
+from monai.metrics.utils import create_ignore_mask
 
 
 class AsymmetricFocalTverskyLoss(_Loss):
@@ -74,19 +75,9 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         if y_true.shape != y_pred.shape:
             raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
 
-        # Build mask after one_hot conversion
-        mask = torch.ones_like(y_true)
-        if self.ignore_index is not None:
-            if original_y_true is not None and self.to_onehot_y:
-                # Use original labels to build spatial mask
-                spatial_mask = (original_y_true != self.ignore_index).float()
-            elif self.ignore_index < y_true.shape[1]:
-                # For already one-hot: use ignored class channel
-                spatial_mask = 1.0 - y_true[:, self.ignore_index : self.ignore_index + 1]
-            else:
-                # For sentinel values: any valid channel
-                spatial_mask = (y_true.sum(dim=1, keepdim=True) > 0).float()
-            mask = spatial_mask.expand_as(y_true)
+        mask = create_ignore_mask(original_y_true if original_y_true is not None else y_true, self.ignore_index)
+        if mask is not None:
+            mask = mask.expand_as(y_true)
             y_pred = y_pred * mask
             y_true = y_true * mask
 
@@ -169,18 +160,12 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         y_pred = torch.clamp(y_pred, self.epsilon, 1.0 - self.epsilon)
         cross_entropy = -y_true * torch.log(y_pred)
 
-        # Build mask from original labels if available
-        spatial_mask: torch.Tensor | None = None
-        if self.ignore_index is not None:
-            if original_y_true is not None and self.to_onehot_y:
-                spatial_mask = (original_y_true != self.ignore_index).float()
-            elif self.ignore_index < y_true.shape[1]:
-                spatial_mask = 1.0 - y_true[:, self.ignore_index : self.ignore_index + 1]
-            else:
-                spatial_mask = (y_true.sum(dim=1, keepdim=True) > 0).float()
-
-            if spatial_mask is not None:
-                cross_entropy = cross_entropy * spatial_mask.expand_as(cross_entropy)
+        # Apply mask from original labels if available
+        mask = create_ignore_mask(original_y_true if original_y_true is not None else y_true, self.ignore_index)
+        if mask is not None:
+            cross_entropy = cross_entropy * mask.expand_as(cross_entropy)
+            # Rename for compatibility with reduction block below
+            spatial_mask = mask
 
         back_ce = torch.pow(1 - y_pred[:, 0], self.gamma) * cross_entropy[:, 0]
         back_ce = (1 - self.delta) * back_ce
@@ -276,30 +261,16 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         if y_pred.shape[1] == 1:
             y_pred = torch.cat([1 - y_pred, y_pred], dim=1)
 
-        # Move one_hot conversion OUTSIDE the if y_pred.shape[1] == 1 block
         if self.to_onehot_y:
-            if self.ignore_index is not None:
-                mask = (y_true != self.ignore_index).float()
-                y_true_clean = torch.where(y_true == self.ignore_index, 0, y_true)
-                y_true = one_hot(y_true_clean, num_classes=self.num_classes)
-                # Keep the channel-wise mask
-                y_true = y_true * mask
-            else:
-                y_true = one_hot(y_true, num_classes=self.num_classes)
-
-        # Check if shapes match
-        if y_true.shape[1] == 1 and y_pred.shape[1] == 2:
-            if self.ignore_index is not None:
-                # Create mask for valid pixels
-                mask = (y_true != self.ignore_index).float()
-                # Set ignore_index values to 0 before conversion
-                y_true_clean = y_true * mask
-                # Convert to 2-channel
-                y_true = torch.cat([1 - y_true_clean, y_true_clean], dim=1)
-                # Apply mask to both channels so ignored pixels are all zeros
-                y_true = y_true * mask
-            else:
-                y_true = torch.cat([1 - y_true, y_true], dim=1)
+            y_true = one_hot(y_true, num_classes=self.num_classes)
+        elif y_true.shape[1] == 1 and y_pred.shape[1] == 2:
+            y_true = torch.cat([1 - y_true, y_true], dim=1)
+
+        original_y_true_unified = y_true  # Use transformed y_true as baseline if original unavailable
+        mask = create_ignore_mask(original_y_true_unified if self.ignore_index is not None else None, self.ignore_index)
+
+        if mask is not None:
+            y_true = y_true * mask
 
         if y_true.shape != y_pred.shape:
             raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")

monai/metrics/confusion_matrix.py

@@ -16,7 +16,7 @@
 
 import torch
 
-from monai.metrics.utils import do_metric_reduction, ignore_background
+from monai.metrics.utils import create_ignore_mask, do_metric_reduction, ignore_background
 from monai.utils import MetricReduction, ensure_tuple
 
 from .metric import CumulativeIterationMetric
@@ -169,14 +169,7 @@ def get_confusion_matrix(
     batch_size, n_class = y_pred.shape[:2]
 
     # Create spatial mask if ignore_index is provided
-    mask = None
-    if ignore_index is not None:
-        if ignore_index >= n_class:
-            # If ignore_index is outside channel range (e.g. 255), we assume it's a spatial mask
-            mask = y.sum(dim=1, keepdim=True) > 0
-        else:
-            # If ignore_index is a valid channel, exclude that specific channel
-            mask = 1.0 - y[:, ignore_index : ignore_index + 1]
+    mask = create_ignore_mask(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)

monai/metrics/generalized_dice.py

@@ -13,7 +13,7 @@
 
 import torch
 
-from monai.metrics.utils import do_metric_reduction
+from monai.metrics.utils import create_ignore_mask, do_metric_reduction
 from monai.utils import MetricReduction, Weight, deprecated_arg, look_up_option
 
 from .metric import CumulativeIterationMetric
@@ -156,13 +156,8 @@ def compute_generalized_dice(
         raise ValueError(f"y_pred - {y_pred.shape} - and y - {y.shape} - should have the same shapes.")
 
     # Apply ignore_index masking
-    if ignore_index is not None:
-        if 0 <= ignore_index < y.shape[1]:
-            # For one-hot: use the ignored class channel
-            mask = 1.0 - y[:, ignore_index : ignore_index + 1]
-        else:
-            # For sentinel values (like 255 or -100), check if any channel is valid
-            mask = (y.sum(dim=1, keepdim=True) > 0).float()
+    mask = create_ignore_mask(y, ignore_index)
+    if mask is not None:
         y_pred = y_pred * mask
         y = y * mask
 

monai/metrics/hausdorff_distance.py

@@ -17,7 +17,13 @@
 import numpy as np
 import torch
 
-from monai.metrics.utils import do_metric_reduction, get_edge_surface_distance, ignore_background, prepare_spacing
+from monai.metrics.utils import (
+    create_ignore_mask,
+    do_metric_reduction,
+    get_edge_surface_distance,
+    ignore_background,
+    prepare_spacing,
+)
 from monai.utils import MetricReduction, convert_data_type
 
 from .metric import CumulativeIterationMetric
@@ -51,7 +57,6 @@ class HausdorffDistanceMetric(CumulativeIterationMetric):
             ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", will not do reduction.
         get_not_nans: whether to return the `not_nans` count, if True, aggregate() returns (metric, not_nans).
             Here `not_nans` count the number of not nans for the metric, thus its shape equals to the shape of the metric.
-        ignore_index: index of the class to ignore during calculation. Defaults to ``None``.
 
     """
 
@@ -100,9 +105,8 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor, **kwargs: Any)
         if dims < 3:
             raise ValueError("y_pred should have at least three dimensions.")
 
-        mask = None
-        if self.ignore_index is not None:
-            mask = (y != self.ignore_index).all(dim=1, keepdim=True).float()
+        mask = create_ignore_mask(y, self.ignore_index)
+        if mask is not None:
             y_pred = y_pred * mask
             y = y * mask
 
@@ -115,8 +119,6 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor, **kwargs: Any)
             percentile=self.percentile,
             directed=self.directed,
             spacing=kwargs.get("spacing"),
-            ignore_index=self.ignore_index,
-            mask=mask,
         )
 
     def aggregate(
@@ -148,8 +150,6 @@ def compute_hausdorff_distance(
     percentile: float | None = None,
     directed: bool = False,
     spacing: int | float | np.ndarray | Sequence[int | float | np.ndarray | Sequence[int | float]] | None = None,
-    mask: torch.Tensor | None = None,
-    ignore_index: int | None = None,
 ) -> torch.Tensor:
     """
     Compute the Hausdorff distance.
@@ -175,7 +175,6 @@ def compute_hausdorff_distance(
             If inner sequence has length 1, isotropic spacing with that value is used for all images in the batch,
             else the inner sequence length must be equal to the image dimensions. If ``None``, spacing of unity is used
             for all images in batch. Defaults to ``None``.
-        ignore_index: index of the class to ignore during calculation. Defaults to ``None``.
     """
 
     if not include_background:
@@ -196,23 +195,12 @@ def compute_hausdorff_distance(
         yp = y_pred[b, c]
         yt = y[b, c]
 
-        if ignore_index is not None:
-            valid_mask = y[b].sum(dim=0) > 0
-            yp = yp * valid_mask
-            yt = yt * valid_mask
-
-            # if everything is ignored, define distance as 0
-            if not valid_mask.any():
-                hd[b, c] = torch.tensor(0.0, device=y_pred.device)
-                continue
-
         _, distances, _ = get_edge_surface_distance(
             yp,
             yt,
             distance_metric=distance_metric,
             spacing=spacing_list[b],
             symmetric=not directed,
-            mask=mask[b, 0] if mask is not None else None,
         )
 
         if len(distances) == 0:

monai/metrics/meandice.py

@@ -13,7 +13,7 @@
 
 import torch
 
-from monai.metrics.utils import do_metric_reduction
+from monai.metrics.utils import create_ignore_mask, do_metric_reduction
 from monai.utils import MetricReduction, deprecated_arg
 
 from .metric import CumulativeIterationMetric
@@ -297,7 +297,7 @@ def compute_channel(self, y_pred: torch.Tensor, y: torch.Tensor, mask: torch.Ten
             mask: binary mask where 0 indicates voxels to ignore.
         """
         if mask is not None:
-            y_pred = y_pred * mask
+            y_pred = y_pred & mask.bool()
             y = y * mask
 
         y_o = torch.sum(y)
@@ -336,19 +336,7 @@ def __call__(self, y_pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor | tupl
             y_pred = y_pred > 0.5
 
         # Create global mask for ignored voxels if ignore_index is set
-        mask = None
-        if self.ignore_index is not None:
-            if y.shape[1] == 1:
-                # Single channel - values are class indices
-                mask = y != self.ignore_index
-            else:
-                # Multi-channel (one-hot or class probabilities)
-                if self.ignore_index < n_pred_ch:
-                    # Class-based ignore: ignore specific class channel
-                    mask = y[:, self.ignore_index : self.ignore_index + 1] == 0
-                else:
-                    # Sentinel-based ignore: ignore where all channels are 0
-                    mask = y.sum(dim=1, keepdim=True) > 0
+        mask = create_ignore_mask(y, self.ignore_index)
 
         first_ch = 0 if self.include_background else 1
         data = []

monai/metrics/meaniou.py

@@ -13,7 +13,7 @@
 
 import torch
 
-from monai.metrics.utils import do_metric_reduction, ignore_background
+from monai.metrics.utils import create_ignore_mask, do_metric_reduction, ignore_background
 from monai.utils import MetricReduction
 
 from .metric import CumulativeIterationMetric
@@ -143,15 +143,9 @@ def compute_iou(
     if y.shape != y_pred.shape:
         raise ValueError(f"y_pred and y should have same shapes, got {y_pred.shape} and {y.shape}.")
 
-    if ignore_index is not None:
-        if ignore_index < y.shape[1]:
-            # For one-hot: mask based on the ignored class channel
-            mask = 1.0 - y[:, ignore_index : ignore_index + 1]
-            if mask.shape != y_pred.shape:
-                mask = mask.expand_as(y_pred)
-        else:
-            # For sentinel values, check if any channel is valid
-            mask = (y.sum(dim=1, keepdim=True) > 0).float()
+    mask = create_ignore_mask(y, ignore_index)
+    if mask is not None:
+        if mask.shape != y_pred.shape:
             mask = mask.expand_as(y_pred)
         y_pred = y_pred * mask
         y = y * mask