Add loss_fn to IgniteMetric and rename to IgniteMetricHandler

monai/handlers/__init__.py

@@ -42,3 +42,4 @@
 from .tensorboard_handlers import TensorBoardHandler, TensorBoardImageHandler, TensorBoardStatsHandler
 from .utils import from_engine, ignore_data, stopping_fn_from_loss, stopping_fn_from_metric, write_metrics_reports
 from .validation_handler import ValidationHandler
+from .dice_ce_metric import DiceCEMetric

monai/handlers/dice_ce_metric.py

@@ -0,0 +1,63 @@
+# 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.
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from monai.handlers.ignite_metric import IgniteMetric
+from monai.utils import MetricReduction
+from monai.losses import DiceCELoss
+from monai.metrics import LossMetric
+
+class DiceCEMetric(IgniteMetric):
+    """
+    Computes DiceCE score metric from full size Tensor and collects average over batch, class-channels, iterations.
+    """
+
+    def __init__(
+        self,
+        # include_background: bool = True,
+        reduction: MetricReduction | str = MetricReduction.MEAN,
+        # num_classes: int | None = None,
+        output_transform: Callable = lambda x: x,
+        save_details: bool = True,
+        *args,
+        **kwargs
+    ) -> None:
+        """
+
+        Args:
+            include_background: whether to include dice computation on the first channel of the predicted output.
+                Defaults to True.
+            reduction: define the mode to reduce metrics, will only execute reduction on `not-nan` values,
+                available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``,
+                ``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", will not do reduction.
+            num_classes: number of input channels (always including the background). When this is None,
+                ``y_pred.shape[1]`` will be used. This option is useful when both ``y_pred`` and ``y`` are
+                single-channel class indices and the number of classes is not automatically inferred from data.
+            output_transform: callable to extract `y_pred` and `y` from `ignite.engine.state.output` then
+                construct `(y_pred, y)` pair, where `y_pred` and `y` can be `batch-first` Tensors or
+                lists of `channel-first` Tensors. the form of `(y_pred, y)` is required by the `update()`.
+                `engine.state` and `output_transform` inherit from the ignite concept:
+                https://pytorch.org/ignite/concepts.html#state, explanation and usage example are in the tutorial:
+                https://github.com/Project-MONAI/tutorials/blob/master/modules/batch_output_transform.ipynb.
+            save_details: whether to save metric computation details per image, for example: mean dice of every image.
+                default to True, will save to `engine.state.metric_details` dict with the metric name as key.
+            args: Arguments for the DiceCELoss
+
+        See also:
+            :py:meth:`monai.metrics.meandice.compute_dice`
+        """
+        # metric_fn = DiceMetric(include_background=include_background, reduction=reduction, num_classes=num_classes)
+        loss_function = DiceCELoss(*args, **kwargs)
+        metric_fn = LossMetric(loss_fn=loss_function, reduction=reduction, get_not_nans=False)
+        super().__init__(metric_fn=metric_fn, output_transform=output_transform, save_details=save_details)

tests/test_handler_dice_ce_metric.py

@@ -0,0 +1,98 @@
+# 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.
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+from ignite.engine import Engine, Events
+from parameterized import parameterized
+
+from monai.handlers import DiceCEMetric, from_engine
+from tests.utils import assert_allclose
+
+TEST_CASE_1 = [{"include_background": True, "output_transform": from_engine(["pred", "label"])}, 0.813259, (4, 2)]
+TEST_CASE_2 = [{"include_background": False, "output_transform": from_engine(["pred", "label"])}, 0.813259, (4, 1)]
+TEST_CASE_3 = [
+    {"reduction": "mean_channel", "output_transform": from_engine(["pred", "label"])},
+    torch.Tensor([0.313262, 2.313251, 0.313262, 0.313262]),
+    (4, 2),
+]
+
+
+class TestHandlerDiceCEMetric(unittest.TestCase):
+    # TODO test multi node averaged dice
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
+    def test_compute(self, input_params, expected_avg, details_shape):
+        dice_metric = DiceCEMetric(**input_params)
+
+        def _val_func(engine, batch):
+            pass
+
+        engine = Engine(_val_func)
+        dice_metric.attach(engine=engine, name="dice_ce_metric")
+        # test input a list of channel-first tensor
+        y_pred = [torch.Tensor([[0], [1]]), torch.Tensor([[1], [0]])]
+        y = torch.Tensor([[[0], [1]], [[0], [1]]])
+        engine.state.output = {"pred": y_pred, "label": y}
+        engine.fire_event(Events.ITERATION_COMPLETED)
+
+        y_pred = [torch.Tensor([[0], [1]]), torch.Tensor([[1], [0]])]
+        y = torch.Tensor([[[0], [1]], [[1], [0]]])
+        engine.state.output = {"pred": y_pred, "label": y}
+        engine.fire_event(Events.ITERATION_COMPLETED)
+
+        engine.fire_event(Events.EPOCH_COMPLETED)
+        assert_allclose(engine.state.metrics["dice_ce_metric"], expected_avg, atol=1e-4, rtol=1e-4, type_test=False)
+        self.assertTupleEqual(tuple(engine.state.metric_details["dice_ce_metric"].shape), details_shape)
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    def test_shape_mismatch(self, input_params, _expected_avg, _details_shape):
+        dice_metric = DiceCEMetric(**input_params)
+        with self.assertRaises((AssertionError, ValueError)):
+            y_pred = torch.Tensor([[0, 1], [1, 0]])
+            y = torch.ones((2, 3))
+            dice_metric.update([y_pred, y])
+
+        with self.assertRaises((AssertionError, ValueError)):
+            y_pred = torch.Tensor([[0, 1], [1, 0]])
+            y = torch.ones((3, 2))
+            dice_metric.update([y_pred, y])
+
+    # @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
+    # def test_compute_n_class(self, input_params, expected_avg, details_shape):
+    #     dice_metric = DiceCEMetric(num_classes=2, **input_params)
+
+    #     def _val_func(engine, batch):
+    #         pass
+
+    #     engine = Engine(_val_func)
+    #     dice_metric.attach(engine=engine, name="dice_ce_metric")
+    #     # test input a list of channel-first tensor
+    #     y_pred = [torch.Tensor([[1]]), torch.Tensor([[0]])]
+    #     y = torch.Tensor([[[0], [1]], [[0], [1]]])
+    #     engine.state.output = {"pred": y_pred, "label": y}
+    #     engine.fire_event(Events.ITERATION_COMPLETED)
+
+    #     y_pred = [torch.Tensor([[1]]), torch.Tensor([[0]])]  # class indices y_pred
+    #     y = torch.Tensor([[[1]], [[0]]])  # class indices y
+    #     engine.state.output = {"pred": y_pred, "label": y}
+    #     engine.fire_event(Events.ITERATION_COMPLETED)
+
+    #     engine.fire_event(Events.EPOCH_COMPLETED)
+    #     assert_allclose(engine.state.metrics["dice_ce_metric"], expected_avg, atol=1e-4, rtol=1e-4, type_test=False)
+    #     self.assertTupleEqual(tuple(engine.state.metric_details["dice_ce_metric"].shape), details_shape)
+
+
+if __name__ == "__main__":
+    unittest.main()