init: starting implementation of ignore_index #8734

monai/check_ignore.py

@@ -0,0 +1,170 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import warnings
+import importlib.util
+from monai.losses import DiceLoss, FocalLoss, AsymmetricUnifiedFocalLoss
+from monai.metrics import (
+    DiceMetric,
+    GeneralizedDiceScore,
+    HausdorffDistanceMetric,
+    SurfaceDistanceMetric,
+    SurfaceDiceMetric,
+)
+from monai.networks import one_hot
+
+
+class IgnoreIndexTester:
+    def __init__(self):
+        self.target_ignore = torch.tensor(
+            [[[[1, 1, 0, 0],
+               [1, 1, 0, 0],
+               [255, 255, 255, 255],
+               [255, 255, 255, 255]]]],
+            dtype=torch.long,
+        )
+
+        self.target_std = torch.tensor(
+            [[[[1, 1, 0, 0],
+               [1, 1, 0, 0],
+               [0, 0, 1, 1],
+               [0, 0, 1, 1]]]],
+            dtype=torch.long,
+        )
+
+        self.input_1 = torch.zeros((1, 1, 4, 4))
+        self.input_1[:, :, 0:2, 0:2] = 5.0
+        self.input_1[:, :, 2:4, 2:4] = -5.0
+
+        self.input_2 = self.input_1.clone()
+        self.input_2[:, :, 3, 3] = 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]
+
+        v1 = torch.as_tensor(v1)
+        v2 = torch.as_tensor(v2)
+
+        match = torch.allclose(v1, v2, atol=1e-4, 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}")
+
+    # ---------------- LOSSES ----------------
+
+    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")
+
+    # ---------------- METRICS ----------------
+
+    def run_metric_test(self, name, metric_cls, **kwargs):
+        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(metric, probs, target):
+            metric.reset()
+
+            if isinstance(
+                metric,
+                (
+                    GeneralizedDiceScore,
+                    HausdorffDistanceMetric,
+                    SurfaceDistanceMetric,
+                    SurfaceDiceMetric,
+                ),
+            ):
+                y_pred = (probs > 0.5).float()
+                t_clean = torch.where(target == 255, 0, target)
+                y = one_hot(t_clean, num_classes=probs.shape[1])
+                y = y * (target != 255).float()
+            else:
+                y_pred = probs.argmax(dim=1, keepdim=True)
+                y = target
+
+            metric(y_pred=y_pred, y=y)
+            return metric.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")
+
+    # ---------------- EXEC ----------------
+
+    def execute(self):
+        print("--- Starting IgnoreIndex Tests (4x4 Geometry) ---")
+
+        self.run_loss_test("DiceLoss", DiceLoss, sigmoid=True)
+        self.run_loss_test("FocalLoss", FocalLoss, use_softmax=False)
+
+        # Unified Focal Loss: single-channel logits only
+        c_ign = AsymmetricUnifiedFocalLoss(ignore_index=255)
+        self.log_result(
+            "UnifiedFocal",
+            c_ign(self.input_1, self.target_ignore),
+            c_ign(self.input_2, self.target_ignore),
+            False,
+            "ignore",
+        )
+
+        c_std = AsymmetricUnifiedFocalLoss()
+        self.log_result(
+            "UnifiedFocal",
+            c_std(self.input_1, self.target_std),
+            c_std(self.input_2, self.target_std),
+            False,
+            "standard",
+        )
+
+        self.run_metric_test("DiceMetric", DiceMetric, include_background=True)
+        self.run_metric_test("GenDice", GeneralizedDiceScore, include_background=True)
+
+        if importlib.util.find_spec("scipy") is not None:
+            self.run_metric_test("Hausdorff", HausdorffDistanceMetric, include_background=True)
+            self.run_metric_test("SurfaceDist", SurfaceDistanceMetric, include_background=True)
+            self.run_metric_test(
+                "SurfaceDice",
+                SurfaceDiceMetric,
+                class_thresholds=[1.0, 1.0],
+                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()