feat: implement ignore_index for losses and metrics

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

Author: Rusheel86

monai/check_ignore.py

@@ -0,0 +1,129 @@
+import torch
+import warnings
+import importlib.util
+from monai.losses import DiceLoss, TverskyLoss, FocalLoss
+from monai.metrics import DiceMetric, MeanIoU, ConfusionMatrixMetric, HausdorffDistanceMetric
+from monai.networks import one_hot
+
+class IgnoreIndexTester:
+    def __init__(self):
+        # Target: [Batch, Channel, H, W]
+        # Bottom row is ignored (255)
+        self.target_ignore = torch.tensor(
+            [[[[1, 0],
+               [255, 255]]]],
+            dtype=torch.long,
+        )
+
+        self.target_std = torch.tensor(
+            [[[[1, 0],
+               [0, 1]]]],
+            dtype=torch.long,
+        )
+
+        self.input_1 = torch.tensor(
+            [[[[2.0, -2.0],
+               [5.0, -5.0]]]]
+        )
+
+        self.input_2 = torch.tensor(
+            [[[[2.0, -2.0],
+               [-5.0, 5.0]]]]
+        )
+
+        self.results = []
+
+    def log_result(self, name, v1, v2, should_match, mode):
+        if isinstance(v1, (list, tuple)):
+            v1 = v1[0]
+        if isinstance(v2, (list, tuple)):
+            v2 = v2[0]
+
+        match = torch.allclose(v1.float(), v2.float(), atol=1e-5, equal_nan=True)
+        success = match == should_match
+        status = "PASS" if success else "FAIL"
+        self.results.append(f"{name:18} [{mode:8}] : {status}")
+
+        if not success:
+            print(f"DEBUG {name} ({mode}) -> {v1.item():.6f} vs {v2.item():.6f}")
+
+    def run_loss_test(self, name, loss_cls, **kwargs):
+        crit_ignore = loss_cls(ignore_index=255, **kwargs)
+        v1_ign = crit_ignore(self.input_1, self.target_ignore)
+        v2_ign = crit_ignore(self.input_2, self.target_ignore)
+        self.log_result(name, v1_ign, v2_ign, True, "ignore")
+
+        crit_std = loss_cls(**kwargs)
+        v1_std = crit_std(self.input_1, self.target_std)
+        v2_std = crit_std(self.input_2, self.target_std)
+        self.log_result(name, v1_std, v2_std, False, "standard")
+
+    def run_metric_test(self, name, metric_cls, **kwargs):
+        # Probabilities for classes 0 and 1
+        p1 = torch.cat([1 - torch.sigmoid(self.input_1), torch.sigmoid(self.input_1)], dim=1)
+        p2 = torch.cat([1 - torch.sigmoid(self.input_2), torch.sigmoid(self.input_2)], dim=1)
+
+        def eval_metric(m, p_probs, t):
+            m.reset()
+            
+            # Hausdorff and distance metrics usually require One-Hot inputs
+            if "Hausdorff" in name:
+                y_pred = (p_probs > 0.5).float()
+                # Clean target for one_hot conversion
+                t_clean = t.clone()
+                mask_ignore = (t == 255)
+                t_clean[mask_ignore] = 0
+                y_true = one_hot(t_clean, num_classes=p_probs.shape[1])
+                # Zero out ignored regions in one-hot ground truth
+                if mask_ignore.any():
+                    y_true[mask_ignore.expand_as(y_true)] = 0
+            else:
+                # Standard overlap metrics use Argmax indices
+                y_pred = p_probs.argmax(dim=1, keepdim=True)
+                y_true = t
+                
+            m(y_pred=y_pred, y=y_true)
+            return m.aggregate()
+
+        m_ignore = metric_cls(ignore_index=255, **kwargs)
+        v1_ign = eval_metric(m_ignore, p1, self.target_ignore)
+        v2_ign = eval_metric(m_ignore, p2, self.target_ignore)
+        self.log_result(name, v1_ign, v2_ign, True, "ignore")
+
+        m_std = metric_cls(**kwargs)
+        v1_std = eval_metric(m_std, p1, self.target_std)
+        v2_std = eval_metric(m_std, p2, self.target_std)
+        self.log_result(name, v1_std, v2_std, False, "standard")
+
+    def run_focal_softmax_standard(self):
+        input_1_mc = torch.cat([self.input_1, -self.input_1], dim=1)
+        input_2_mc = torch.cat([self.input_2, -self.input_2], dim=1)
+
+        crit = FocalLoss(use_softmax=True, to_onehot_y=True)
+        v1 = crit(input_1_mc, self.target_std)
+        v2 = crit(input_2_mc, self.target_std)
+        self.log_result("FocalLoss", v1, v2, False, "softmax")
+
+    def execute(self):
+        print("--- Starting IgnoreIndex Tests ---")
+
+        self.run_loss_test("DiceLoss", DiceLoss, sigmoid=True)
+        self.run_loss_test("TverskyLoss", TverskyLoss, sigmoid=True)
+        self.run_loss_test("FocalLoss", FocalLoss, use_softmax=False)
+        self.run_focal_softmax_standard()
+
+        self.run_metric_test("DiceMetric", DiceMetric, include_background=True)
+        self.run_metric_test("MeanIoU", MeanIoU, include_background=True)
+        self.run_metric_test("Accuracy", ConfusionMatrixMetric, metric_name="accuracy")
+        
+        if importlib.util.find_spec("scipy") is not None:
+            self.run_metric_test("Hausdorff", HausdorffDistanceMetric, include_background=True)
+
+        print("\n--- TEST SUMMARY ---")
+        for r in self.results:
+            print(r)
+
+if __name__ == "__main__":
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        IgnoreIndexTester().execute()