Repository: Project-MONAI/MONAI

Issue:
Binary prediction bug in `HoVerNetNuclearTypePostProcessing`
**Describe the bug**
Now in `HoVerNetNuclearTypePostProcessing`, if set `return_binary` to True, it will change `inst_pred` in place which will affect instance map prediction.
https://github.com/Project-MONAI/MONAI/blob/72130602340847e15f2042b2bb1b64f6f46eeb93/monai/apps/pathology/transforms/post/dictionary.py#L543-L545
Also, users may not only want an instance information dictionary but a type prediction map.

**Expected behavior**
When setting `return_binary` to True, instance map prediction won't be changed.
`HoVerNetNuclearTypePostProcessing` can also output type map prediction when `output_classes` is not None.


Discussion:

Patch:
diff --git a/monai/apps/pathology/transforms/post/array.py b/monai/apps/pathology/transforms/post/array.py
--- a/monai/apps/pathology/transforms/post/array.py
+++ b/monai/apps/pathology/transforms/post/array.py
@@ -12,6 +12,7 @@
 from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
+import torch
 
 from monai.config.type_definitions import DtypeLike, NdarrayOrTensor
 from monai.transforms import Activations, AsDiscrete, BoundingRect, RemoveSmallObjects, SobelGradients
@@ -629,8 +630,8 @@ def __call__(  # type: ignore
 
 class HoVerNetNuclearTypePostProcessing(Transform):
     """
-    The whole post-procesing transform for nuclear type classification branch. It return a dict which contains
-    centroid, bounding box, type prediciton for each instance.
+    The whole post-procesing transform for nuclear type classification branch. It returns type prediction map and a
+    dict contains centroid, bounding box, type prediciton for each instance.
 
     Args:
         min_num_points: assumed that the created contour does not form a contour if it does not contain more points
@@ -656,7 +657,9 @@ def __init__(
         self.generate_instance_centroid = GenerateInstanceCentroid()
         self.generate_instance_type = GenerateInstanceType()
 
-    def __call__(self, type_pred: NdarrayOrTensor, instance_pred: NdarrayOrTensor) -> Dict:  # type: ignore
+    def __call__(  # type: ignore
+        self, type_pred: NdarrayOrTensor, instance_pred: NdarrayOrTensor
+    ) -> Union[Tuple[NdarrayOrTensor, Dict], Dict]:
         type_pred = Activations(softmax=True)(type_pred)
         type_pred = AsDiscrete(argmax=True)(type_pred)
 
@@ -691,4 +694,11 @@ def __call__(self, type_pred: NdarrayOrTensor, instance_pred: NdarrayOrTensor) -
                 inst_info_dict[inst_id]["type"] = inst_type  # type: ignore
                 inst_info_dict[inst_id]["type_probability"] = type_prob  # type: ignore
 
+            instance_pred = convert_to_dst_type(instance_pred, type_pred)[0]
+            pred_type_map = torch.zeros_like(instance_pred)  # type: ignore
+            for key, value in inst_info_dict.items():
+                pred_type_map[instance_pred == key] = value["type"]  # type: ignore
+
+            return pred_type_map, inst_info_dict
+
         return inst_info_dict
diff --git a/monai/apps/pathology/transforms/post/dictionary.py b/monai/apps/pathology/transforms/post/dictionary.py
--- a/monai/apps/pathology/transforms/post/dictionary.py
+++ b/monai/apps/pathology/transforms/post/dictionary.py
@@ -523,6 +523,7 @@ def __init__(
         self.converter = HoVerNetNuclearTypePostProcessing(
             min_num_points=min_num_points, level=level, return_centroids=return_centroids, output_classes=output_classes
         )
+        self.output_classes = output_classes
         self.type_pred_key = type_pred_key
         self.instance_pred_key = instance_pred_key
         self.pred_binary_key = pred_binary_key
@@ -533,16 +534,21 @@ def __call__(self, data):
         d = dict(data)
         inst_pred = d[self.instance_pred_key]
         type_pred = d[self.type_pred_key]
-        inst_info_dict = self.converter(type_pred, inst_pred)
+        if self.output_classes is not None:
+            pred_type_map, inst_info_dict = self.converter(type_pred, inst_pred)
+            d[self.type_pred_key] = pred_type_map
+        else:
+            inst_info_dict = self.converter(type_pred, inst_pred)
+
         key_to_add = f"{self.instance_info_dict_key}"
         if key_to_add in d:
             raise KeyError(f"Type information with key {key_to_add} already exists.")
         d[key_to_add] = inst_info_dict
 
-        inst_pred = convert_to_dst_type(inst_pred, type_pred)[0]
         if self.return_binary:
-            inst_pred[inst_pred > 0] = 1
-            d[self.pred_binary_key] = inst_pred
+            d[self.pred_binary_key] = (inst_pred > 0).astype(int)
+        inst_pred = convert_to_dst_type(inst_pred, type_pred)[0]
+        d[self.instance_pred_key] = inst_pred
 
         return d
 


Test patch:
diff --git a/tests/test_hovernet_nuclear_type_post_processing.py b/tests/test_hovernet_nuclear_type_post_processing.py
--- a/tests/test_hovernet_nuclear_type_post_processing.py
+++ b/tests/test_hovernet_nuclear_type_post_processing.py
@@ -74,7 +74,7 @@ def test_value(self, in_type, test_data, seg_postpprocessing, kwargs, expected):
         post_transforms = HoVerNetNuclearTypePostProcessing(**kwargs)
         out = post_transforms(type_pred=in_type(test_data[None]), instance_pred=pred["dist"])
         if out is not None:
-            assert_allclose(out[1]["centroid"], expected[1], type_test=False)
+            assert_allclose(out[1][1]["centroid"], expected[1], type_test=False)
 
 
 if __name__ == "__main__":