From 9379b13b78f46f7e9bc05bf12d50b6f665f0f698 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Mon, 29 Mar 2021 14:12:28 -0400
Subject: [PATCH 1/9] Implement ProbMapGenerator handler

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 monai/apps/pathology/handlers.py | 99 ++++++++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)
 create mode 100644 monai/apps/pathology/handlers.py

diff --git a/monai/apps/pathology/handlers.py b/monai/apps/pathology/handlers.py
new file mode 100644
index 0000000000..054314502e
--- /dev/null
+++ b/monai/apps/pathology/handlers.py
@@ -0,0 +1,99 @@
+import logging
+import os
+from typing import TYPE_CHECKING, Dict, Optional
+
+import numpy as np
+
+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 ProbMapGenerator:
+    """
+    Event handler triggered on completing every iteration to save the probability map
+    """
+
+    def __init__(
+        self,
+        output_dir: str = "./",
+        name: Optional[str] = None,
+    ) -> None:
+        """
+        Args:
+            output_dir: output directory to save probability maps.
+            name: identifier of logging.logger to use, defaulting to `engine.logger`.
+
+        """
+        self.logger = logging.getLogger(name)
+        self._name = name
+        self.output_dir = output_dir
+        self.prob_map: Dict[str, np.ndarray] = {}
+        self.level: Dict[str, int] = {}
+        self.counter: Dict[str, int] = {}
+        self.num_done_images: int = 0
+        self.num_images: int = 0
+
+    def attach(self, engine: Engine) -> None:
+        """
+        Args:
+            engine: Ignite Engine, it can be a trainer, validator or evaluator.
+        """
+
+        self.num_images = len(engine.data_loader.dataset.data_list)
+
+        for sample in engine.data_loader.dataset.data_list:
+            name = sample["name"]
+            self.prob_map[name] = np.zeros(sample["mask_shape"])
+            self.counter[name] = len(sample["mask_locations"])
+            self.level[name] = sample["level"]
+
+        if self._name is None:
+            self.logger = engine.logger
+        if not engine.has_event_handler(self, Events.ITERATION_COMPLETED):
+            engine.add_event_handler(Events.ITERATION_COMPLETED, self)
+        if not engine.has_event_handler(self.finalize, Events.COMPLETED):
+            engine.add_event_handler(Events.COMPLETED, self.finalize)
+
+    def __call__(self, engine: Engine) -> None:
+        """
+        This method assumes self.batch_transform will extract metadata from the input batch.
+
+        Args:
+            engine: Ignite Engine, it can be a trainer, validator or evaluator.
+        """
+        names = engine.state.batch["name"]
+        locs = engine.state.batch["mask_location"]
+        pred = engine.state.output["pred"]
+        for i, name in enumerate(names):
+            self.prob_map[name][locs[0][i], locs[1][i]] = pred[i]
+            self.counter[name] -= 1
+            if self.counter[name] == 0:
+                self.save_prob_map(name)
+
+    def save_prob_map(self, name: str) -> None:
+        """
+        This method save the probability map for an image, when its inference is finished,
+        and delete that probability map from memory.
+
+        Args:
+            name: the name of image to be saved.
+        """
+        file_path = os.path.join(self.output_dir, name)
+        np.save(file_path + ".npy", self.prob_map[name])
+
+        self.num_done_images += 1
+        self.logger.info(f"Inference of '{name}' is done [{self.num_done_images}/{self.num_images}]!")
+        del self.prob_map[name]
+        del self.counter[name]
+        del self.level[name]
+
+    def finalize(self, engine: Engine):
+        if self.counter:
+            raise RuntimeError(f"Counter: {self.counter}")
+        else:
+            self.logger.info(f"Probability map is created for {self.num_done_images}/{self.num_images} images!")

From e15cf72d58ed046d6b8c34cdb6de9f7ee2a752ee Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Mon, 29 Mar 2021 14:12:48 -0400
Subject: [PATCH 2/9] Update doc and init

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 docs/source/apps.rst             | 4 ++++
 monai/apps/pathology/__init__.py | 1 +
 2 files changed, 5 insertions(+)

diff --git a/docs/source/apps.rst b/docs/source/apps.rst
index 0c92d4c443..de7c40f3a7 100644
--- a/docs/source/apps.rst
+++ b/docs/source/apps.rst
@@ -75,3 +75,7 @@ Applications
 .. automodule:: monai.apps.pathology.utils
 .. autoclass:: PathologyProbNMS
     :members:
+
+.. automodule:: monai.apps.pathology.handlers
+.. autoclass:: ProbMapGenerator
+    :members:
\ No newline at end of file
diff --git a/monai/apps/pathology/__init__.py b/monai/apps/pathology/__init__.py
index 3af25365ba..ce2c3b5340 100644
--- a/monai/apps/pathology/__init__.py
+++ b/monai/apps/pathology/__init__.py
@@ -11,3 +11,4 @@
 
 from .datasets import PatchWSIDataset, SmartCacheDataset
 from .utils import ProbNMS
+from .handlers import ProbMapGenerator

From adfff4e54aaab1f8dd0d904780a13a766c5443ed Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Mon, 29 Mar 2021 16:01:06 -0400
Subject: [PATCH 3/9] Sort init imports

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 monai/apps/pathology/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/monai/apps/pathology/__init__.py b/monai/apps/pathology/__init__.py
index ce2c3b5340..a5f105e303 100644
--- a/monai/apps/pathology/__init__.py
+++ b/monai/apps/pathology/__init__.py
@@ -10,5 +10,5 @@
 # limitations under the License.
 
 from .datasets import PatchWSIDataset, SmartCacheDataset
-from .utils import ProbNMS
 from .handlers import ProbMapGenerator
+from .utils import ProbNMS

From d75c6572ae2bbe84b93f49e65d15e4601657e78b Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Mon, 29 Mar 2021 21:05:23 -0400
Subject: [PATCH 4/9] Add unittest for ProbMapGenerator

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/test_handler_prob_map_generator.py | 95 ++++++++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 tests/test_handler_prob_map_generator.py

diff --git a/tests/test_handler_prob_map_generator.py b/tests/test_handler_prob_map_generator.py
new file mode 100644
index 0000000000..2702f8f167
--- /dev/null
+++ b/tests/test_handler_prob_map_generator.py
@@ -0,0 +1,95 @@
+# 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 os
+import unittest
+
+import numpy as np
+import torch
+from ignite.engine import Engine
+from parameterized import parameterized
+from torch.utils.data import DataLoader
+
+from monai.apps.pathology.handlers import ProbMapGenerator
+from monai.data.dataset import Dataset
+from monai.engines import Evaluator
+from monai.handlers import ValidationHandler
+
+TEST_CASE_0 = ["image_inference_output_1", 2]
+TEST_CASE_1 = ["image_inference_output_2", 9]
+TEST_CASE_2 = ["image_inference_output_3", 1000]
+
+
+class TestDataset(Dataset):
+    def __init__(self, name, size):
+        self.data_list = [
+            {
+                "name": name,
+                "mask_shape": (size, size),
+                "mask_locations": [[i, i] for i in range(size)],
+                "level": 0,
+            }
+        ]
+        self.len = size
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, index):
+        return {
+            "name": self.data_list[0]["name"],
+            "mask_location": self.data_list[0]["mask_locations"][index],
+            "pred": index + 1,
+        }
+
+
+class TestEvaluator(Evaluator):
+    def _iteration(self, engine, batchdata):
+        return batchdata
+
+
+class TestHandlerProbMapGenerator(unittest.TestCase):
+    @parameterized.expand(
+        [
+            TEST_CASE_0,
+            TEST_CASE_1,
+            TEST_CASE_2,
+        ]
+    )
+    def test_prob_map_generator(self, name, size):
+        # set up dataset
+        dataset = TestDataset(name, size)
+        data_loader = DataLoader(dataset, batch_size=1)
+
+        # set up engine
+        def inference(enging, batch):
+            pass
+
+        engine = Engine(inference)
+
+        # add ProbMapGenerator() to evaluator
+        output_dir = os.path.join(os.path.dirname(__file__), "testing_data")
+        prob_map_gen = ProbMapGenerator(output_dir=output_dir)
+
+        evaluator = TestEvaluator(torch.device("cpu:0"), data_loader, size, val_handlers=[prob_map_gen])
+
+        # set up validation handler
+        validation = ValidationHandler(evaluator, interval=1)
+        validation.attach(engine)
+
+        engine.run(data_loader)
+
+        prob_map = np.load(os.path.join(output_dir, name + ".npy"))
+        self.assertListEqual(np.diag(prob_map).astype(int).tolist(), list(range(1, size + 1)))
+
+
+if __name__ == "__main__":
+    unittest.main()

From 55c1c86e183262801a0af55213860db4fdeb8aca Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Mon, 29 Mar 2021 21:45:17 -0400
Subject: [PATCH 5/9] Ignore if ignite is not available

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/test_handler_prob_map_generator.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/test_handler_prob_map_generator.py b/tests/test_handler_prob_map_generator.py
index 2702f8f167..783c22c822 100644
--- a/tests/test_handler_prob_map_generator.py
+++ b/tests/test_handler_prob_map_generator.py
@@ -11,10 +11,10 @@
 
 import os
 import unittest
+from unittest import skipUnless
 
 import numpy as np
 import torch
-from ignite.engine import Engine
 from parameterized import parameterized
 from torch.utils.data import DataLoader
 
@@ -23,6 +23,8 @@
 from monai.engines import Evaluator
 from monai.handlers import ValidationHandler
 
+Events, has_ignite = optional_import("ignite.engine", "0.4.4", exact_version, "Events")
+
 TEST_CASE_0 = ["image_inference_output_1", 2]
 TEST_CASE_1 = ["image_inference_output_2", 9]
 TEST_CASE_2 = ["image_inference_output_3", 1000]
@@ -64,6 +66,7 @@ class TestHandlerProbMapGenerator(unittest.TestCase):
             TEST_CASE_2,
         ]
     )
+    @skipUnless(has_ignite, "Requires pytorch-ignite.")
     def test_prob_map_generator(self, name, size):
         # set up dataset
         dataset = TestDataset(name, size)

From eee4f60c7152ddbd219a0d84096137909d608b0b Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Mon, 29 Mar 2021 22:46:55 -0400
Subject: [PATCH 6/9] Update Engine import

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/test_handler_prob_map_generator.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/test_handler_prob_map_generator.py b/tests/test_handler_prob_map_generator.py
index 783c22c822..098242baf0 100644
--- a/tests/test_handler_prob_map_generator.py
+++ b/tests/test_handler_prob_map_generator.py
@@ -22,8 +22,9 @@
 from monai.data.dataset import Dataset
 from monai.engines import Evaluator
 from monai.handlers import ValidationHandler
+from monai.utils import exact_version, optional_import
 
-Events, has_ignite = optional_import("ignite.engine", "0.4.4", exact_version, "Events")
+Engine, has_ignite = optional_import("ignite.engine", "0.4.4", exact_version, "Engine")
 
 TEST_CASE_0 = ["image_inference_output_1", 2]
 TEST_CASE_1 = ["image_inference_output_2", 9]

From ad3138146a8467d8c06f7878717906190754b7b6 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Mon, 29 Mar 2021 23:19:58 -0400
Subject: [PATCH 7/9] Exclude from min-test

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 tests/min_tests.py                       | 1 +
 tests/test_handler_prob_map_generator.py | 6 +-----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/tests/min_tests.py b/tests/min_tests.py
index 83c1ceea9f..e896e81c70 100644
--- a/tests/min_tests.py
+++ b/tests/min_tests.py
@@ -43,6 +43,7 @@ def run_testsuit():
         "test_handler_confusion_matrix_dist",
         "test_handler_hausdorff_distance",
         "test_handler_mean_dice",
+        "test_handler_prob_map_generator",
         "test_handler_rocauc",
         "test_handler_rocauc_dist",
         "test_handler_segmentation_saver",
diff --git a/tests/test_handler_prob_map_generator.py b/tests/test_handler_prob_map_generator.py
index 098242baf0..2702f8f167 100644
--- a/tests/test_handler_prob_map_generator.py
+++ b/tests/test_handler_prob_map_generator.py
@@ -11,10 +11,10 @@
 
 import os
 import unittest
-from unittest import skipUnless
 
 import numpy as np
 import torch
+from ignite.engine import Engine
 from parameterized import parameterized
 from torch.utils.data import DataLoader
 
@@ -22,9 +22,6 @@
 from monai.data.dataset import Dataset
 from monai.engines import Evaluator
 from monai.handlers import ValidationHandler
-from monai.utils import exact_version, optional_import
-
-Engine, has_ignite = optional_import("ignite.engine", "0.4.4", exact_version, "Engine")
 
 TEST_CASE_0 = ["image_inference_output_1", 2]
 TEST_CASE_1 = ["image_inference_output_2", 9]
@@ -67,7 +64,6 @@ class TestHandlerProbMapGenerator(unittest.TestCase):
             TEST_CASE_2,
         ]
     )
-    @skipUnless(has_ignite, "Requires pytorch-ignite.")
     def test_prob_map_generator(self, name, size):
         # set up dataset
         dataset = TestDataset(name, size)

From bbfc0c893ed5547364b8367fa6cb13573be1e52a Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 30 Mar 2021 15:36:34 -0400
Subject: [PATCH 8/9] Address all the comments

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 docs/source/apps.rst                     |  2 +-
 monai/apps/pathology/__init__.py         |  2 +-
 monai/apps/pathology/handlers.py         | 16 +++++++++-------
 tests/test_handler_prob_map_generator.py | 10 +++++-----
 4 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/docs/source/apps.rst b/docs/source/apps.rst
index de7c40f3a7..215942a804 100644
--- a/docs/source/apps.rst
+++ b/docs/source/apps.rst
@@ -77,5 +77,5 @@ Applications
     :members:
 
 .. automodule:: monai.apps.pathology.handlers
-.. autoclass:: ProbMapGenerator
+.. autoclass:: ProbMapProducer
     :members:
\ No newline at end of file
diff --git a/monai/apps/pathology/__init__.py b/monai/apps/pathology/__init__.py
index a5f105e303..2751ca7f5b 100644
--- a/monai/apps/pathology/__init__.py
+++ b/monai/apps/pathology/__init__.py
@@ -10,5 +10,5 @@
 # limitations under the License.
 
 from .datasets import PatchWSIDataset, SmartCacheDataset
-from .handlers import ProbMapGenerator
+from .handlers import ProbMapProducer
 from .utils import ProbNMS
diff --git a/monai/apps/pathology/handlers.py b/monai/apps/pathology/handlers.py
index 054314502e..2f395c030f 100644
--- a/monai/apps/pathology/handlers.py
+++ b/monai/apps/pathology/handlers.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 
+from monai.config import DtypeLike
 from monai.utils import exact_version, optional_import
 
 Events, _ = optional_import("ignite.engine", "0.4.4", exact_version, "Events")
@@ -13,7 +14,7 @@
     Engine, _ = optional_import("ignite.engine", "0.4.4", exact_version, "Engine")
 
 
-class ProbMapGenerator:
+class ProbMapProducer:
     """
     Event handler triggered on completing every iteration to save the probability map
     """
@@ -21,11 +22,15 @@ class ProbMapGenerator:
     def __init__(
         self,
         output_dir: str = "./",
+        output_postfix: str = "",
+        dtype: DtypeLike = np.float64,
         name: Optional[str] = None,
     ) -> None:
         """
         Args:
             output_dir: output directory to save probability maps.
+            output_postfix: a string appended to all output file names.
+            dtype: the data type in which the probability map is stored. Default np.float64.
             name: identifier of logging.logger to use, defaulting to `engine.logger`.
 
         """
@@ -44,9 +49,9 @@ def attach(self, engine: Engine) -> None:
             engine: Ignite Engine, it can be a trainer, validator or evaluator.
         """
 
-        self.num_images = len(engine.data_loader.dataset.data_list)
+        self.num_images = len(engine.data_loader.dataset.data)
 
-        for sample in engine.data_loader.dataset.data_list:
+        for sample in engine.data_loader.dataset.data:
             name = sample["name"]
             self.prob_map[name] = np.zeros(sample["mask_shape"])
             self.counter[name] = len(sample["mask_locations"])
@@ -93,7 +98,4 @@ def save_prob_map(self, name: str) -> None:
         del self.level[name]
 
     def finalize(self, engine: Engine):
-        if self.counter:
-            raise RuntimeError(f"Counter: {self.counter}")
-        else:
-            self.logger.info(f"Probability map is created for {self.num_done_images}/{self.num_images} images!")
+        self.logger.info(f"Probability map is created for {self.num_done_images}/{self.num_images} images!")
diff --git a/tests/test_handler_prob_map_generator.py b/tests/test_handler_prob_map_generator.py
index 2702f8f167..4882060be9 100644
--- a/tests/test_handler_prob_map_generator.py
+++ b/tests/test_handler_prob_map_generator.py
@@ -18,7 +18,7 @@
 from parameterized import parameterized
 from torch.utils.data import DataLoader
 
-from monai.apps.pathology.handlers import ProbMapGenerator
+from monai.apps.pathology.handlers import ProbMapProducer
 from monai.data.dataset import Dataset
 from monai.engines import Evaluator
 from monai.handlers import ValidationHandler
@@ -30,7 +30,7 @@
 
 class TestDataset(Dataset):
     def __init__(self, name, size):
-        self.data_list = [
+        self.data = [
             {
                 "name": name,
                 "mask_shape": (size, size),
@@ -45,8 +45,8 @@ def __len__(self):
 
     def __getitem__(self, index):
         return {
-            "name": self.data_list[0]["name"],
-            "mask_location": self.data_list[0]["mask_locations"][index],
+            "name": self.data[0]["name"],
+            "mask_location": self.data[0]["mask_locations"][index],
             "pred": index + 1,
         }
 
@@ -77,7 +77,7 @@ def inference(enging, batch):
 
         # add ProbMapGenerator() to evaluator
         output_dir = os.path.join(os.path.dirname(__file__), "testing_data")
-        prob_map_gen = ProbMapGenerator(output_dir=output_dir)
+        prob_map_gen = ProbMapProducer(output_dir=output_dir)
 
         evaluator = TestEvaluator(torch.device("cpu:0"), data_loader, size, val_handlers=[prob_map_gen])
 

From 193cf5c51f5748ac05581beef9641f093911a311 Mon Sep 17 00:00:00 2001
From: Behrooz <3968947+behxyz@users.noreply.github.com>
Date: Tue, 30 Mar 2021 17:55:56 -0400
Subject: [PATCH 9/9] Fix file path and dtype

Signed-off-by: Behrooz <3968947+behxyz@users.noreply.github.com>
---
 monai/apps/pathology/handlers.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/monai/apps/pathology/handlers.py b/monai/apps/pathology/handlers.py
index 2f395c030f..046e403e0f 100644
--- a/monai/apps/pathology/handlers.py
+++ b/monai/apps/pathology/handlers.py
@@ -37,6 +37,8 @@ def __init__(
         self.logger = logging.getLogger(name)
         self._name = name
         self.output_dir = output_dir
+        self.output_postfix = output_postfix
+        self.dtype = dtype
         self.prob_map: Dict[str, np.ndarray] = {}
         self.level: Dict[str, int] = {}
         self.counter: Dict[str, int] = {}
@@ -53,7 +55,7 @@ def attach(self, engine: Engine) -> None:
 
         for sample in engine.data_loader.dataset.data:
             name = sample["name"]
-            self.prob_map[name] = np.zeros(sample["mask_shape"])
+            self.prob_map[name] = np.zeros(sample["mask_shape"], dtype=self.dtype)
             self.counter[name] = len(sample["mask_locations"])
             self.level[name] = sample["level"]
 
@@ -89,7 +91,7 @@ def save_prob_map(self, name: str) -> None:
             name: the name of image to be saved.
         """
         file_path = os.path.join(self.output_dir, name)
-        np.save(file_path + ".npy", self.prob_map[name])
+        np.save(file_path + self.output_postfix + ".npy", self.prob_map[name])
 
         self.num_done_images += 1
         self.logger.info(f"Inference of '{name}' is done [{self.num_done_images}/{self.num_images}]!")