feat: implement ignore_index support for losses and metrics

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

Author: Rusheel86

monai/losses/unified_focal_loss.py

@@ -262,12 +262,22 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
 
         # Check if shapes match
         if y_true.shape[1] == 1 and y_pred.shape[1] == 2:
-            y_true = torch.cat([1 - y_true, y_true], dim=1)
+            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)
+
         if y_true.shape != y_pred.shape:
             raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
-
-        if torch.max(y_true) != self.num_classes - 1:
-            raise ValueError(f"Please make sure the number of classes is {self.num_classes - 1}")
+        if self.ignore_index is None and torch.max(y_true) > self.num_classes - 1:
+            raise ValueError(f"Invalid class index found. Maximum class should be {self.num_classes - 1}")
 
         asy_focal_loss = self.asy_focal_loss(y_pred, y_true)
         asy_focal_tversky_loss = self.asy_focal_tversky_loss(y_pred, y_true)

monai/metrics/generalized_dice.py

@@ -154,6 +154,12 @@ def compute_generalized_dice(
     if y.shape != y_pred.shape:
         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:
+        mask = (y != ignore_index).all(dim=1, keepdim=True).float()
+        y_pred = y_pred * mask
+        y = y * mask
+
     n_channels = y_pred.shape[1]
     channels_to_use = list(range(n_channels))
 
@@ -188,17 +194,17 @@ def compute_generalized_dice(
     else:
         w_full = torch.ones_like(y_o_float)
 
-    w = w_full[:, channels_to_use]
-
     # Replace infinite values for non-appearing classes by the maximum weight
-    for b_idx in range(w.shape[0]):
-        batch_w = w[b_idx]
+    for b_idx in range(w_full.shape[0]):
+        batch_w = w_full[b_idx]
         infs = torch.isinf(batch_w)
         if infs.any():
             batch_w[infs] = 0
             max_w = torch.max(batch_w)
             batch_w[infs] = max_w if max_w > 0 else 1.0
 
+    w = w_full[:, channels_to_use]
+
     if sum_over_classes:
         intersection = (intersection * w).sum(dim=1, keepdim=True)
         denominator = (denominator * w).sum(dim=1, keepdim=True)

monai/metrics/surface_dice.py

@@ -277,7 +277,16 @@ def compute_surface_dice(
                 distances_gt_pred <= class_thresholds[c]
             )
         else:
-            areas_pred, areas_gt = areas  # type: ignore
+            # Handle areas being returned as a single item or a tuple
+            if isinstance(areas, (list, tuple)):
+                if len(areas) == 2:
+                    areas_pred, areas_gt = areas
+                elif len(areas) == 1:
+                    areas_pred = areas_gt = areas[0]
+                else:
+                    areas_pred = areas_gt = torch.tensor([], device=y_pred.device)
+            else:
+                areas_pred = areas_gt = areas
             areas_gt, areas_pred = areas_gt[edges_gt], areas_pred[edges_pred]
             boundary_complete = areas_gt.sum() + areas_pred.sum()
             gt_true = areas_gt[distances_gt_pred <= class_thresholds[c]].sum() if len(areas_gt) > 0 else 0.0

monai/metrics/utils.py

@@ -374,7 +374,15 @@ def get_edge_surface_distance(
 
     distances = tuple(d if d is not None else edges_pred.new_empty((0,)) for d in distances)
 
-    areas = edge_results[3:] if use_subvoxels else ()
+    areas = edge_results[2:] if use_subvoxels else ()
+
+    # Ensure areas is always a tuple of 2 when use_subvoxels=True
+    if use_subvoxels and isinstance(areas, (list, tuple)):
+        if len(areas) == 1:
+            areas = (areas[0], areas[0])
+        elif len(areas) != 2:
+            # Unexpected length, create empty tensors
+            areas = (torch.tensor([], device=y_pred.device), torch.tensor([], device=y_pred.device))
 
     out = convert_to_tensor(((edges_pred, edges_gt), distances, tuple(areas)), device=y_pred.device)  # type: ignore[no-any-return]
 

tests/metrics/test_ignore_index_metrics.py

@@ -47,7 +47,6 @@
 
 @unittest.skipUnless(has_scipy, "Scipy required for surface metrics")
 class TestIgnoreIndexMetrics(unittest.TestCase):
-
     @parameterized.expand(TEST_METRICS + SCIPY_METRICS)
     def test_metric_ignore_consistency(self, metric_class, kwargs):
         # Initialize metric with ignore_index