2223 Add PostProcessing handler

docs/source/handlers.rst

@@ -155,6 +155,11 @@ Transform inverter
 .. autoclass:: TransformInverter
     :members:
 
+Post processing
+---------------
+.. autoclass:: PostProcessing
+    :members:
+
 Utilities
 ---------
 .. automodule:: monai.handlers.utils

monai/handlers/__init__.py

@@ -22,6 +22,7 @@
 from .metric_logger import MetricLogger, MetricLoggerKeys
 from .metrics_saver import MetricsSaver
 from .parameter_scheduler import ParamSchedulerHandler
+from .post_processing import PostProcessing
 from .regression_metrics import MeanAbsoluteError, MeanSquaredError, PeakSignalToNoiseRatio, RootMeanSquaredError
 from .roc_auc import ROCAUC
 from .segmentation_saver import SegmentationSaver

monai/handlers/post_processing.py

@@ -0,0 +1,56 @@
+# Copyright 2020 - 2021 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 typing import TYPE_CHECKING, Callable
+
+from monai.engines.utils import IterationEvents, engine_apply_transform
+from monai.utils import exact_version, optional_import
+
+Events, _ = optional_import("ignite.engine", "0.4.4", exact_version, "Events")
+if TYPE_CHECKING:
+    from ignite.engine import Engine
+else:
+    Engine, _ = optional_import("ignite.engine", "0.4.4", exact_version, "Engine")
+
+
+class PostProcessing:
+    """
+    Ignite handler to execute additional post processing after the post transforms in engines.
+    So users can insert other handlers between post transforms and this post processing handler.
+
+    """
+
+    def __init__(self, transform: Callable) -> None:
+        """
+        Args:
+            transform: callable function to execute on the `engine.state.batch` and `engine.state.output`.
+                can also be composed post transforms.
+
+        """
+        self.transform = transform
+
+    def attach(self, engine: Engine) -> None:
+        """
+        Args:
+            engine: Ignite Engine, it can be a trainer, validator or evaluator.
+        """
+        engine.add_event_handler(IterationEvents.MODEL_COMPLETED, self)
+
+    def __call__(self, engine: Engine) -> None:
+        """
+        Args:
+            engine: Ignite Engine, it can be a trainer, validator or evaluator.
+        """
+        engine.state.batch, engine.state.output = engine_apply_transform(
+            batch=engine.state.batch,
+            output=engine.state.output,
+            transform=self.transform,
+        )

tests/min_tests.py

@@ -124,6 +124,7 @@ def run_testsuit():
         "test_testtimeaugmentation",
         "test_cachedataset_persistent_workers",
         "test_invertd",
+        "test_handler_post_processing",
     ]
     assert sorted(exclude_cases) == sorted(set(exclude_cases)), f"Duplicated items in {exclude_cases}"
 

tests/test_handler_post_processing.py

@@ -0,0 +1,62 @@
+# Copyright 2020 - 2021 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 unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.engines import SupervisedEvaluator
+from monai.handlers import PostProcessing
+from monai.transforms import Activationsd, AsDiscreted, Compose, CopyItemsd
+
+# test lambda function as `transform`
+TEST_CASE_1 = [{"transform": lambda x: dict(pred=x["pred"] + 1.0)}, torch.tensor([[[[1.9975], [1.9997]]]])]
+# test composed post transforms as `transform`
+TEST_CASE_2 = [
+    {
+        "transform": Compose(
+            [
+                CopyItemsd(keys="filename", times=1, names="filename_bak"),
+                AsDiscreted(keys="pred", threshold_values=True, to_onehot=True, n_classes=2),
+            ]
+        )
+    },
+    torch.tensor([[[[1.0], [1.0]], [[0.0], [0.0]]]]),
+]
+
+
+class TestHandlerPostProcessing(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    def test_compute(self, input_params, expected):
+        data = [
+            {"image": torch.tensor([[[[2.0], [3.0]]]]), "filename": "test1"},
+            {"image": torch.tensor([[[[6.0], [8.0]]]]), "filename": "test2"},
+        ]
+        # set up engine, PostProcessing handler works together with post_transform of engine
+        engine = SupervisedEvaluator(
+            device=torch.device("cpu:0"),
+            val_data_loader=data,
+            epoch_length=2,
+            network=torch.nn.PReLU(),
+            post_transform=Compose([Activationsd(keys="pred", sigmoid=True)]),
+            val_handlers=[PostProcessing(**input_params)],
+        )
+        engine.run()
+
+        torch.testing.assert_allclose(engine.state.output["pred"], expected)
+        filename = engine.state.output.get("filename_bak")
+        if filename is not None:
+            self.assertEqual(filename, "test2")
+
+
+if __name__ == "__main__":
+    unittest.main()