Repository: Project-MONAI/MONAI

Issue:
MetaTensor support for `ProbMapProducer`



Discussion:

Patch:
diff --git a/monai/apps/auto3dseg/hpo_gen.py b/monai/apps/auto3dseg/hpo_gen.py
--- a/monai/apps/auto3dseg/hpo_gen.py
+++ b/monai/apps/auto3dseg/hpo_gen.py
@@ -39,7 +39,6 @@ def get_hyperparameters(self):
 
     @abstractmethod
     def update_params(self, *args, **kwargs):
-
         """Update Algo parameters according to the hyperparameters to be evaluated."""
         raise NotImplementedError
 
@@ -50,7 +49,6 @@ def set_score(self):
 
     @abstractmethod
     def run_algo(self, *args, **kwargs):
-
         """Interface for launch the training given the fetched hyperparameters."""
         raise NotImplementedError
 
diff --git a/monai/handlers/probability_maps.py b/monai/handlers/probability_maps.py
--- a/monai/handlers/probability_maps.py
+++ b/monai/handlers/probability_maps.py
@@ -18,6 +18,7 @@
 
 from monai.config import DtypeLike, IgniteInfo
 from monai.utils import ProbMapKeys, min_version, optional_import
+from monai.utils.enums import CommonKeys
 
 Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events")
 if TYPE_CHECKING:
@@ -28,7 +29,10 @@
 
 class ProbMapProducer:
     """
-    Event handler triggered on completing every iteration to calculate and save the probability map
+    Event handler triggered on completing every iteration to calculate and save the probability map.
+    This handler use metadata from MetaTensor to create the probability map. This can be simply achieved by using
+    `monai.data.SlidingPatchWSIDataset` or `monai.data.MaskedPatchWSIDataset` as the dataset.
+
     """
 
     def __init__(
@@ -91,8 +95,8 @@ def __call__(self, engine: Engine) -> None:
         """
         if not isinstance(engine.state.batch, dict) or not isinstance(engine.state.output, dict):
             raise ValueError("engine.state.batch and engine.state.output must be dictionaries.")
-        names = engine.state.batch["metadata"][ProbMapKeys.NAME]
-        locs = engine.state.batch["metadata"][ProbMapKeys.LOCATION]
+        names = engine.state.batch[CommonKeys.IMAGE].meta[ProbMapKeys.NAME]
+        locs = engine.state.batch[CommonKeys.IMAGE].meta[ProbMapKeys.LOCATION]
         probs = engine.state.output[self.prob_key]
         for name, loc, prob in zip(names, locs, probs):
             self.prob_map[name][tuple(loc)] = prob


Test patch:
diff --git a/tests/test_handler_prob_map_producer.py b/tests/test_handler_prob_map_producer.py
--- a/tests/test_handler_prob_map_producer.py
+++ b/tests/test_handler_prob_map_producer.py
@@ -16,14 +16,13 @@
 import torch
 from ignite.engine import Engine
 from parameterized import parameterized
-from torch.utils.data import DataLoader
 
-from monai.data.dataset import Dataset
+from monai.data import DataLoader, Dataset, MetaTensor
 from monai.engines import Evaluator
 from monai.handlers import ProbMapProducer, ValidationHandler
 from monai.utils.enums import ProbMapKeys
 
-TEST_CASE_0 = ["temp_image_inference_output_1", 2]
+TEST_CASE_0 = ["temp_image_inference_output_1", 1]
 TEST_CASE_1 = ["temp_image_inference_output_2", 9]
 TEST_CASE_2 = ["temp_image_inference_output_3", 100]
 
@@ -35,32 +34,32 @@ def __init__(self, name, size):
                 {
                     "image": name,
                     ProbMapKeys.COUNT.value: size,
-                    ProbMapKeys.SIZE.value: np.array([size, size]),
+                    ProbMapKeys.SIZE.value: np.array([size + 1, size + 1]),
                     ProbMapKeys.LOCATION.value: np.array([i, i + 1]),
                 }
-                for i in range(size - 1)
+                for i in range(size)
             ]
         )
         self.image_data = [
             {
                 ProbMapKeys.NAME.value: name,
                 ProbMapKeys.COUNT.value: size,
-                ProbMapKeys.SIZE.value: np.array([size, size]),
+                ProbMapKeys.SIZE.value: np.array([size + 1, size + 1]),
             }
         ]
 
     def __getitem__(self, index):
-        return {
-            "image": np.zeros((3, 2, 2)),
+
+        image = np.ones((3, 2, 2)) * index
+        metadata = {
             ProbMapKeys.COUNT.value: self.data[index][ProbMapKeys.COUNT.value],
-            "metadata": {
-                ProbMapKeys.NAME.value: self.data[index]["image"],
-                ProbMapKeys.SIZE.value: self.data[index][ProbMapKeys.SIZE.value],
-                ProbMapKeys.LOCATION.value: self.data[index][ProbMapKeys.LOCATION.value],
-            },
-            "pred": index + 1,
+            ProbMapKeys.NAME.value: self.data[index]["image"],
+            ProbMapKeys.SIZE.value: self.data[index][ProbMapKeys.SIZE.value],
+            ProbMapKeys.LOCATION.value: self.data[index][ProbMapKeys.LOCATION.value],
         }
 
+        return {"image": MetaTensor(x=image, meta=metadata), "pred": index + 1}
+
 
 class TestEvaluator(Evaluator):
     def _iteration(self, engine, batchdata):
@@ -72,10 +71,11 @@ class TestHandlerProbMapGenerator(unittest.TestCase):
     def test_prob_map_generator(self, name, size):
         # set up dataset
         dataset = TestDataset(name, size)
-        data_loader = DataLoader(dataset, batch_size=1)
+        batch_size = 2
+        data_loader = DataLoader(dataset, batch_size=batch_size)
 
         # set up engine
-        def inference(enging, batch):
+        def inference(engine, batch):
             pass
 
         engine = Engine(inference)
@@ -84,7 +84,9 @@ def inference(enging, batch):
         output_dir = os.path.join(os.path.dirname(__file__), "testing_data")
         prob_map_gen = ProbMapProducer(output_dir=output_dir)
 
-        evaluator = TestEvaluator(torch.device("cpu:0"), data_loader, size, val_handlers=[prob_map_gen])
+        evaluator = TestEvaluator(
+            torch.device("cpu:0"), data_loader, np.ceil(size / batch_size), val_handlers=[prob_map_gen]
+        )
 
         # set up validation handler
         validation = ValidationHandler(interval=1, validator=None)
@@ -94,8 +96,8 @@ def inference(enging, batch):
         engine.run(data_loader)
 
         prob_map = np.load(os.path.join(output_dir, name + ".npy"))
-        self.assertListEqual(np.vstack(prob_map.nonzero()).T.tolist(), [[i, i + 1] for i in range(size - 1)])
-        self.assertListEqual(prob_map[prob_map.nonzero()].tolist(), [i + 1 for i in range(size - 1)])
+        self.assertListEqual(np.vstack(prob_map.nonzero()).T.tolist(), [[i, i + 1] for i in range(size)])
+        self.assertListEqual(prob_map[prob_map.nonzero()].tolist(), [i + 1 for i in range(size)])
 
 
 if __name__ == "__main__":