diff --git a/docs/source/apps.rst b/docs/source/apps.rst index 4c45a5fb39..0c92d4c443 100644 --- a/docs/source/apps.rst +++ b/docs/source/apps.rst @@ -71,3 +71,7 @@ Applications :members: .. autoclass:: SmartCachePatchWSIDataset :members: + +.. automodule:: monai.apps.pathology.utils +.. autoclass:: PathologyProbNMS + :members: diff --git a/docs/source/utils.rst b/docs/source/utils.rst index e0b993da60..071d9ecefd 100644 --- a/docs/source/utils.rst +++ b/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 diff --git a/monai/apps/pathology/__init__.py b/monai/apps/pathology/__init__.py index bbdb812c03..3af25365ba 100644 --- a/monai/apps/pathology/__init__.py +++ b/monai/apps/pathology/__init__.py @@ -10,3 +10,4 @@ # limitations under the License. from .datasets import PatchWSIDataset, SmartCacheDataset +from .utils import ProbNMS diff --git a/monai/apps/pathology/utils.py b/monai/apps/pathology/utils.py new file mode 100644 index 0000000000..b0803526fd --- /dev/null +++ b/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 diff --git a/monai/utils/__init__.py b/monai/utils/__init__.py index d622ce96ae..f6a137f47d 100644 --- a/monai/utils/__init__.py +++ b/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 diff --git a/monai/utils/prob_nms.py b/monai/utils/prob_nms.py new file mode 100644 index 0000000000..bdffdfe005 --- /dev/null +++ b/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 diff --git a/tests/test_pathology_prob_nms.py b/tests/test_pathology_prob_nms.py new file mode 100644 index 0000000000..223b136ea7 --- /dev/null +++ b/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() diff --git a/tests/test_prob_nms.py b/tests/test_prob_nms.py new file mode 100644 index 0000000000..fb88d9cfb4 --- /dev/null +++ b/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()