1822 Implement Prob NMS

docs/source/apps.rst

@@ -71,3 +71,7 @@ Applications
     :members:
 .. autoclass:: SmartCachePatchWSIDataset
     :members:
+
+.. automodule:: monai.apps.pathology.utils
+.. autoclass:: PathologyProbNMS
+    :members:

docs/source/utils.rst

@@ -27,6 +27,12 @@ Misc
 .. automodule:: monai.utils.misc
   :members:
 
+Prob NMS
+--------
+.. automodule:: monai.utils.prob_nms
+.. autoclass:: ProbNMS
+  :members:
+
 Profiling
 ---------
 .. automodule:: monai.utils.profiling

monai/apps/pathology/__init__.py

@@ -10,3 +10,4 @@
 # limitations under the License.
 
 from .datasets import PatchWSIDataset, SmartCacheDataset
+from .utils import ProbNMS

monai/apps/pathology/utils.py

@@ -0,0 +1,43 @@
+# 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 Union
+
+import numpy as np
+import torch
+
+from monai.utils import ProbNMS
+
+
+class PathologyProbNMS(ProbNMS):
+    """
+    This class extends monai.utils.ProbNMS and add the `resolution` option for
+    Pathology.
+    """
+
+    def __call__(
+        self,
+        probs_map: Union[np.ndarray, torch.Tensor],
+        resolution_level: int = 0,
+    ):
+        """
+        probs_map: the input probabilities map, it must have shape (H[, W, ...]).
+        resolution_level: the level at which the probabilities map is made.
+        """
+        resolution = pow(2, resolution_level)
+        org_outputs = ProbNMS.__call__(self, probs_map)
+        outputs = []
+        for org_output in org_outputs:
+            prob = org_output[0]
+            coord = np.asarray(org_output[1:])
+            coord_wsi = ((coord + 0.5) * resolution).astype(int)
+            outputs.append([prob] + list(coord_wsi))
+        return outputs

monai/utils/__init__.py

@@ -69,5 +69,6 @@
     min_version,
     optional_import,
 )
+from .prob_nms import ProbNMS
 from .profiling import PerfContext, torch_profiler_full, torch_profiler_time_cpu_gpu, torch_profiler_time_end_to_end
 from .state_cacher import StateCacher

monai/utils/prob_nms.py

@@ -0,0 +1,99 @@
+from typing import List, Sequence, Tuple, Union
+
+import numpy as np
+import torch
+
+from monai.networks.layers import GaussianFilter
+
+
+class ProbNMS:
+    """
+    Performs probability based non-maximum suppression (NMS) on the probabilities map via
+    iteratively selecting the coordinate with highest probability and then move it as well
+    as its surrounding values. The remove range is determined by the parameter `box_size`.
+    If multiple coordinates have the same highest probability, only one of them will be
+    selected.
+
+    Args:
+        spatial_dims: number of spatial dimensions of the input probabilities map.
+            Defaults to 2.
+        sigma: the standard deviation for gaussian filter.
+            It could be a single value, or `spatial_dims` number of values. Defaults to 0.0.
+        prob_threshold: the probability threshold, the function will stop searching if
+            the highest probability is no larger than the threshold. The value should be
+            no less than 0.0. Defaults to 0.5.
+        box_size: determines the sizes of the removing area of the selected coordinates for
+            each dimensions. Defaults to 48.
+
+    Return:
+        a list of selected lists, where inner lists contain probability and coordinates.
+        For example, for 3D input, the inner lists are in the form of [probability, x, y, z].
+
+    Raises:
+        ValueError: When ``prob_threshold`` is less than 0.0.
+        ValueError: When ``box_size`` is a list or tuple, and its length is not equal to `spatial_dims`.
+        ValueError: When ``box_size`` has a less than 1 value.
+
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int = 2,
+        sigma: Union[Sequence[float], float, Sequence[torch.Tensor], torch.Tensor] = 0.0,
+        prob_threshold: float = 0.5,
+        box_size: Union[int, List[int], Tuple[int]] = 48,
+    ) -> None:
+        self.sigma = sigma
+        self.spatial_dims = spatial_dims
+        if self.sigma != 0:
+            self.filter = GaussianFilter(spatial_dims=spatial_dims, sigma=sigma)
+        if prob_threshold < 0:
+            raise ValueError("prob_threshold should be no less than 0.0.")
+        self.prob_threshold = prob_threshold
+        if isinstance(box_size, int):
+            self.box_size = np.asarray([box_size] * spatial_dims)
+        else:
+            if len(box_size) != spatial_dims:
+                raise ValueError("the sequence length of box_size should be the same as spatial_dims.")
+            self.box_size = np.asarray(box_size)
+        if self.box_size.min() <= 0:
+            raise ValueError("box_size should be larger than 0.")
+
+        self.box_lower_bd = self.box_size // 2
+        self.box_upper_bd = self.box_size - self.box_lower_bd
+
+    def __call__(
+        self,
+        probs_map: Union[np.ndarray, torch.Tensor],
+    ):
+        """
+        probs_map: the input probabilities map, it must have shape (H[, W, ...]).
+        """
+        if self.sigma != 0:
+            if not isinstance(probs_map, torch.Tensor):
+                probs_map = torch.as_tensor(probs_map, dtype=torch.float)
+            self.filter.to(probs_map)
+            probs_map = self.filter(probs_map)
+        else:
+            if not isinstance(probs_map, torch.Tensor):
+                probs_map = probs_map.copy()
+
+        if isinstance(probs_map, torch.Tensor):
+            probs_map = probs_map.detach().cpu().numpy()
+
+        probs_map_shape = probs_map.shape
+
+        outputs = []
+        while np.max(probs_map) > self.prob_threshold:
+            max_idx = np.unravel_index(probs_map.argmax(), probs_map_shape)
+            prob_max = probs_map[max_idx]
+            max_idx_arr = np.asarray(max_idx)
+            outputs.append([prob_max] + list(max_idx_arr))
+
+            idx_min_range = (max_idx_arr - self.box_lower_bd).clip(0, None)
+            idx_max_range = (max_idx_arr + self.box_upper_bd).clip(None, probs_map_shape)
+            # for each dimension, set values during index ranges to 0
+            slices = tuple([slice(idx_min_range[i], idx_max_range[i]) for i in range(self.spatial_dims)])
+            probs_map[slices] = 0
+
+        return outputs

tests/test_pathology_prob_nms.py

@@ -0,0 +1,57 @@
+# 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 numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.apps.pathology.utils import PathologyProbNMS
+
+probs_map_2d = np.random.rand(100, 100).clip(0, 0.5)
+probs_map_2d[33, 33] = 0.7
+probs_map_2d[66, 66] = 0.9
+expected_2d = [[0.9, 133, 133], [0.7, 67, 67]]
+TEST_CASES_2D = [
+    {"spatial_dims": 2, "prob_threshold": 0.5, "box_size": [10, 10]},
+    {"resolution_level": 1},
+    probs_map_2d,
+    expected_2d,
+]
+
+probs_map_3d = torch.rand([50, 50, 50]).uniform_(0, 0.5)
+probs_map_3d[25, 25, 25] = 0.7
+probs_map_3d[45, 45, 45] = 0.9
+expected_3d = [[0.9, 91, 91, 91], [0.7, 51, 51, 51]]
+TEST_CASES_3D = [
+    {"spatial_dims": 3, "prob_threshold": 0.5, "box_size": (10, 10, 10)},
+    {"resolution_level": 1},
+    probs_map_3d,
+    expected_3d,
+]
+
+
+class TestPathologyProbNMS(unittest.TestCase):
+    @parameterized.expand(
+        [
+            TEST_CASES_2D,
+            TEST_CASES_3D,
+        ]
+    )
+    def test_output(self, class_args, call_args, probs_map, expected):
+        nms = PathologyProbNMS(**class_args)
+        output = nms(probs_map, **call_args)
+        np.testing.assert_allclose(output, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_prob_nms.py

@@ -0,0 +1,103 @@
+# 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 numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.utils import ProbNMS
+
+probs_map_1 = np.random.rand(100, 100).clip(0, 0.5)
+TEST_CASES_2D_1 = [{"spatial_dims": 2, "prob_threshold": 0.5, "box_size": 10}, probs_map_1, []]
+
+probs_map_2 = np.random.rand(100, 100).clip(0, 0.5)
+probs_map_2[33, 33] = 0.7
+probs_map_2[66, 66] = 0.9
+expected_2 = [[0.9, 66, 66], [0.7, 33, 33]]
+TEST_CASES_2D_2 = [
+    {"spatial_dims": 2, "prob_threshold": 0.5, "box_size": [10, 10]},
+    probs_map_2,
+    expected_2,
+]
+
+probs_map_3 = np.random.rand(100, 100).clip(0, 0.5)
+probs_map_3[56, 58] = 0.7
+probs_map_3[60, 66] = 0.8
+probs_map_3[66, 66] = 0.9
+expected_3 = [[0.9, 66, 66], [0.8, 60, 66]]
+TEST_CASES_2D_3 = [
+    {"spatial_dims": 2, "prob_threshold": 0.5, "box_size": (10, 20)},
+    probs_map_3,
+    expected_3,
+]
+
+probs_map_4 = np.random.rand(100, 100).clip(0, 0.5)
+probs_map_4[33, 33] = 0.7
+probs_map_4[66, 66] = 0.9
+expected_4 = [[0.9, 66, 66]]
+TEST_CASES_2D_4 = [
+    {"spatial_dims": 2, "prob_threshold": 0.8, "box_size": 10},
+    probs_map_4,
+    expected_4,
+]
+
+probs_map_5 = np.random.rand(100, 100).clip(0, 0.5)
+TEST_CASES_2D_5 = [{"spatial_dims": 2, "prob_threshold": 0.5, "sigma": 0.1}, probs_map_5, []]
+
+probs_map_6 = torch.as_tensor(np.random.rand(100, 100).clip(0, 0.5))
+TEST_CASES_2D_6 = [{"spatial_dims": 2, "prob_threshold": 0.5, "sigma": 0.1}, probs_map_6, []]
+
+probs_map_7 = torch.as_tensor(np.random.rand(100, 100).clip(0, 0.5))
+probs_map_7[33, 33] = 0.7
+probs_map_7[66, 66] = 0.9
+if torch.cuda.is_available():
+    probs_map_7 = probs_map_7.cuda()
+expected_7 = [[0.9, 66, 66], [0.7, 33, 33]]
+TEST_CASES_2D_7 = [
+    {"spatial_dims": 2, "prob_threshold": 0.5, "sigma": 0.1},
+    probs_map_7,
+    expected_7,
+]
+
+probs_map_3d = torch.rand([50, 50, 50]).uniform_(0, 0.5)
+probs_map_3d[25, 25, 25] = 0.7
+probs_map_3d[45, 45, 45] = 0.9
+expected_3d = [[0.9, 45, 45, 45], [0.7, 25, 25, 25]]
+TEST_CASES_3D = [
+    {"spatial_dims": 3, "prob_threshold": 0.5, "box_size": (10, 10, 10)},
+    probs_map_3d,
+    expected_3d,
+]
+
+
+class TestProbNMS(unittest.TestCase):
+    @parameterized.expand(
+        [
+            TEST_CASES_2D_1,
+            TEST_CASES_2D_2,
+            TEST_CASES_2D_3,
+            TEST_CASES_2D_4,
+            TEST_CASES_2D_5,
+            TEST_CASES_2D_6,
+            TEST_CASES_2D_7,
+            TEST_CASES_3D,
+        ]
+    )
+    def test_output(self, class_args, probs_map, expected):
+        nms = ProbNMS(**class_args)
+        output = nms(probs_map)
+        np.testing.assert_allclose(output, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()