From df7089b25e18fbe617d35eb11b1a56a2c9f147cb Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Wed, 24 Mar 2021 22:46:40 +0800 Subject: [PATCH 1/7] Optimize speed and support any dimensions Signed-off-by: Yiheng Wang --- monai/utils/radial_nms.py | 56 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 monai/utils/radial_nms.py diff --git a/monai/utils/radial_nms.py b/monai/utils/radial_nms.py new file mode 100644 index 0000000000..4efbadb6d5 --- /dev/null +++ b/monai/utils/radial_nms.py @@ -0,0 +1,56 @@ +from typing import Union + +import numpy as np +import torch + +from monai.networks.layers import GaussianFilter + + +class RadialNMS: + + def __init__(self, spatial_dims: int, sigma: float = 0.0, prob_threshold: float = 0.5, radius: int = 24): + self.sigma = sigma + self.spatial_dims = spatial_dims + if self.sigma > 0: + self.filter = GaussianFilter(spatial_dims=spatial_dims, sigma=sigma) + self.prob_threshold = prob_threshold + self.radius = radius + + def __call__(self, probs_map: Union[np.ndarray, torch.Tensor], level: int): + if self.sigma > 0: + if isinstance(probs_map, np.ndarray): + probs_map = torch.as_tensor(probs_map, dtype=torch.float) + device = ( + torch.device("cuda") + if (probs_map.device == "cuda" and torch.cuda.is_available()) + else torch.device("cpu:0") + ) + self.filter.to(device) + probs_map = self.filter(probs_map) + probs_map = probs_map.detach().cpu().numpy() + + probs_map_shape = probs_map.shape + resolution = pow(2, level) + + 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) + coord_wsi = ((max_idx_arr + 0.5) * resolution).astype(int) + outputs.append([prob_max] + list(coord_wsi)) + + # achieve min index for each dimension + idx_min_range = max_idx_arr - self.radius + idx_min_range[idx_min_range < 0] = 0 + + # achieve max index for each dimension + idx_upper_bound = np.asarray(probs_map_shape) + + idx_max_range = max_idx_arr + self.radius + idx_max_range[idx_max_range > idx_upper_bound] = idx_upper_bound[idx_max_range > idx_upper_bound] + # set values during index ranges for each dimension 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 From 65048ac5f1756cf7a3e0823db907c05a8725b761 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Wed, 24 Mar 2021 22:55:12 +0800 Subject: [PATCH 2/7] Fix black issue Signed-off-by: Yiheng Wang --- monai/utils/radial_nms.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monai/utils/radial_nms.py b/monai/utils/radial_nms.py index 4efbadb6d5..47dcc96438 100644 --- a/monai/utils/radial_nms.py +++ b/monai/utils/radial_nms.py @@ -7,7 +7,6 @@ class RadialNMS: - def __init__(self, spatial_dims: int, sigma: float = 0.0, prob_threshold: float = 0.5, radius: int = 24): self.sigma = sigma self.spatial_dims = spatial_dims From 0fe17c4c1d5b25fbd5f8d5d557e30a3f6a8e330f Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Thu, 25 Mar 2021 17:11:40 +0800 Subject: [PATCH 3/7] Add unittest and docstrings Signed-off-by: Yiheng Wang --- docs/source/utils.rst | 5 ++ monai/utils/__init__.py | 1 + monai/utils/pixel_nms.py | 83 +++++++++++++++++++++++++++++ monai/utils/radial_nms.py | 55 -------------------- tests/test_pixel_nms.py | 107 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 196 insertions(+), 55 deletions(-) create mode 100644 monai/utils/pixel_nms.py delete mode 100644 monai/utils/radial_nms.py create mode 100644 tests/test_pixel_nms.py diff --git a/docs/source/utils.rst b/docs/source/utils.rst index e0b993da60..947c6a4b3c 100644 --- a/docs/source/utils.rst +++ b/docs/source/utils.rst @@ -27,6 +27,11 @@ Misc .. automodule:: monai.utils.misc :members: +Pixel NMS +--------- +.. autoclass:: monai.utils.PixelNMS + :members: + Profiling --------- .. automodule:: monai.utils.profiling diff --git a/monai/utils/__init__.py b/monai/utils/__init__.py index d622ce96ae..0c244b4e21 100644 --- a/monai/utils/__init__.py +++ b/monai/utils/__init__.py @@ -69,5 +69,6 @@ min_version, optional_import, ) +from .pixel_nms import PixelNMS 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/pixel_nms.py b/monai/utils/pixel_nms.py new file mode 100644 index 0000000000..51f49324c8 --- /dev/null +++ b/monai/utils/pixel_nms.py @@ -0,0 +1,83 @@ +from typing import Sequence, Union + +import numpy as np +import torch + +from monai.networks.layers import GaussianFilter + + +class PixelNMS: + """ + Performs pixel 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 `pixel_dis`. + 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. + sigma: the standard deviation for gaussian filter. + It could be a single value, or `spatial_dims` number of values. Defaults to 0. + prob_threshold: the probability threshold, the PixelNMS will stop searching if + the highest probability is no larger than the threshold. Defaults to 0.5. + pixel_dis: the distance in pixels that determines the surrounding area of the selected + coordinate that will be removed before the next iteration. The area is a square (cube) + for 2D (3D) input, and the edge length is 2 * pixel_dis. Defaults to 24. + + 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]. + + """ + + def __init__( + self, + spatial_dims: int, + sigma: Union[Sequence[float], float, Sequence[torch.Tensor], torch.Tensor] = 0, + prob_threshold: float = 0.5, + pixel_dis: int = 24, + ) -> None: + self.sigma = sigma + self.spatial_dims = spatial_dims + if self.sigma != 0: + self.filter = GaussianFilter(spatial_dims=spatial_dims, sigma=sigma) + self.prob_threshold = prob_threshold + self.pixel_dis = pixel_dis + + 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 original input is made. The returned + coordinates will be converted according to this value. + """ + 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) + + if isinstance(probs_map, torch.Tensor): + probs_map = probs_map.detach().cpu().numpy() + + probs_map_shape = probs_map.shape + resolution = pow(2, resolution_level) + + 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) + coord_wsi = ((max_idx_arr + 0.5) * resolution).astype(int) + outputs.append([prob_max] + list(coord_wsi)) + + idx_min_range = (max_idx_arr - self.pixel_dis).clip(0, None) + idx_max_range = (max_idx_arr + self.pixel_dis).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/monai/utils/radial_nms.py b/monai/utils/radial_nms.py deleted file mode 100644 index 47dcc96438..0000000000 --- a/monai/utils/radial_nms.py +++ /dev/null @@ -1,55 +0,0 @@ -from typing import Union - -import numpy as np -import torch - -from monai.networks.layers import GaussianFilter - - -class RadialNMS: - def __init__(self, spatial_dims: int, sigma: float = 0.0, prob_threshold: float = 0.5, radius: int = 24): - self.sigma = sigma - self.spatial_dims = spatial_dims - if self.sigma > 0: - self.filter = GaussianFilter(spatial_dims=spatial_dims, sigma=sigma) - self.prob_threshold = prob_threshold - self.radius = radius - - def __call__(self, probs_map: Union[np.ndarray, torch.Tensor], level: int): - if self.sigma > 0: - if isinstance(probs_map, np.ndarray): - probs_map = torch.as_tensor(probs_map, dtype=torch.float) - device = ( - torch.device("cuda") - if (probs_map.device == "cuda" and torch.cuda.is_available()) - else torch.device("cpu:0") - ) - self.filter.to(device) - probs_map = self.filter(probs_map) - probs_map = probs_map.detach().cpu().numpy() - - probs_map_shape = probs_map.shape - resolution = pow(2, level) - - 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) - coord_wsi = ((max_idx_arr + 0.5) * resolution).astype(int) - outputs.append([prob_max] + list(coord_wsi)) - - # achieve min index for each dimension - idx_min_range = max_idx_arr - self.radius - idx_min_range[idx_min_range < 0] = 0 - - # achieve max index for each dimension - idx_upper_bound = np.asarray(probs_map_shape) - - idx_max_range = max_idx_arr + self.radius - idx_max_range[idx_max_range > idx_upper_bound] = idx_upper_bound[idx_max_range > idx_upper_bound] - # set values during index ranges for each dimension 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_pixel_nms.py b/tests/test_pixel_nms.py new file mode 100644 index 0000000000..29a29823f1 --- /dev/null +++ b/tests/test_pixel_nms.py @@ -0,0 +1,107 @@ +# 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 PixelNMS + +probs_map_1 = np.random.rand(100, 100).clip(0, 0.5) +TEST_CASES_2D_1 = [{"spatial_dims": 2, "prob_threshold": 0.5, "pixel_dis": 5}, {"resolution_level": 0}, 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, 133, 133], [0.7, 67, 67]] +TEST_CASES_2D_2 = [ + {"spatial_dims": 2, "prob_threshold": 0.5, "pixel_dis": 5}, + {"resolution_level": 1}, + probs_map_2, + expected_2, +] + +probs_map_3 = np.random.rand(100, 100).clip(0, 0.5) +probs_map_3[33, 33] = 0.7 +probs_map_3[66, 66] = 0.9 +expected_3 = [[0.9, 266, 266]] +TEST_CASES_2D_3 = [ + {"spatial_dims": 2, "prob_threshold": 0.5, "pixel_dis": 40}, + {"resolution_level": 2}, + 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, 266, 266]] +TEST_CASES_2D_4 = [ + {"spatial_dims": 2, "prob_threshold": 0.8, "pixel_dis": 5}, + {"resolution_level": 2}, + 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}, {"resolution_level": 2}, 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}, {"resolution_level": 2}, 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, 266, 266], [0.7, 134, 134]] +TEST_CASES_2D_7 = [ + {"spatial_dims": 2, "prob_threshold": 0.5, "sigma": 0.1}, + {"resolution_level": 2}, + 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, "pixel_dis": 5}, + {"resolution_level": 0}, + probs_map_3d, + expected_3d, +] + + +class TestPixelNMS(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, call_args, probs_map, expected): + nms = PixelNMS(**class_args) + output = nms(probs_map, **call_args) + np.testing.assert_allclose(output, expected) + + +if __name__ == "__main__": + unittest.main() From f0bfb3b42a7bfa540e58d3de7ca4463ec4d10c3d Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Fri, 26 Mar 2021 12:53:52 +0800 Subject: [PATCH 4/7] Modify box size and others Signed-off-by: Yiheng Wang --- docs/source/utils.rst | 6 +- monai/utils/__init__.py | 2 +- monai/utils/{pixel_nms.py => prob_nms.py} | 60 ++++++++++++------- tests/{test_pixel_nms.py => test_prob_nms.py} | 38 ++++++------ 4 files changed, 59 insertions(+), 47 deletions(-) rename monai/utils/{pixel_nms.py => prob_nms.py} (55%) rename tests/{test_pixel_nms.py => test_prob_nms.py} (70%) diff --git a/docs/source/utils.rst b/docs/source/utils.rst index 947c6a4b3c..428160d54d 100644 --- a/docs/source/utils.rst +++ b/docs/source/utils.rst @@ -27,9 +27,9 @@ Misc .. automodule:: monai.utils.misc :members: -Pixel NMS ---------- -.. autoclass:: monai.utils.PixelNMS +Prob NMS +-------- +.. autoclass:: monai.utils.ProbNMS :members: Profiling diff --git a/monai/utils/__init__.py b/monai/utils/__init__.py index 0c244b4e21..f6a137f47d 100644 --- a/monai/utils/__init__.py +++ b/monai/utils/__init__.py @@ -69,6 +69,6 @@ min_version, optional_import, ) -from .pixel_nms import PixelNMS +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/pixel_nms.py b/monai/utils/prob_nms.py similarity index 55% rename from monai/utils/pixel_nms.py rename to monai/utils/prob_nms.py index 51f49324c8..bdffdfe005 100644 --- a/monai/utils/pixel_nms.py +++ b/monai/utils/prob_nms.py @@ -1,4 +1,4 @@ -from typing import Sequence, Union +from typing import List, Sequence, Tuple, Union import numpy as np import torch @@ -6,76 +6,92 @@ from monai.networks.layers import GaussianFilter -class PixelNMS: +class ProbNMS: """ - Performs pixel based non-maximum suppression (NMS) on the probabilities map via + 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 `pixel_dis`. + 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. - prob_threshold: the probability threshold, the PixelNMS will stop searching if - the highest probability is no larger than the threshold. Defaults to 0.5. - pixel_dis: the distance in pixels that determines the surrounding area of the selected - coordinate that will be removed before the next iteration. The area is a square (cube) - for 2D (3D) input, and the edge length is 2 * pixel_dis. Defaults to 24. + 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, - sigma: Union[Sequence[float], float, Sequence[torch.Tensor], torch.Tensor] = 0, + spatial_dims: int = 2, + sigma: Union[Sequence[float], float, Sequence[torch.Tensor], torch.Tensor] = 0.0, prob_threshold: float = 0.5, - pixel_dis: int = 24, + 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 - self.pixel_dis = pixel_dis + 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], - resolution_level: int = 0, ): """ probs_map: the input probabilities map, it must have shape (H[, W, ...]). - resolution_level: the level at which the original input is made. The returned - coordinates will be converted according to this value. """ 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 - resolution = pow(2, resolution_level) 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) - coord_wsi = ((max_idx_arr + 0.5) * resolution).astype(int) - outputs.append([prob_max] + list(coord_wsi)) + outputs.append([prob_max] + list(max_idx_arr)) - idx_min_range = (max_idx_arr - self.pixel_dis).clip(0, None) - idx_max_range = (max_idx_arr + self.pixel_dis).clip(None, probs_map_shape) + 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 diff --git a/tests/test_pixel_nms.py b/tests/test_prob_nms.py similarity index 70% rename from tests/test_pixel_nms.py rename to tests/test_prob_nms.py index 29a29823f1..ffccf62e2a 100644 --- a/tests/test_pixel_nms.py +++ b/tests/test_prob_nms.py @@ -15,29 +15,28 @@ import torch from parameterized import parameterized -from monai.utils import PixelNMS +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, "pixel_dis": 5}, {"resolution_level": 0}, probs_map_1, []] +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, 133, 133], [0.7, 67, 67]] +expected_2 = [[0.9, 66, 66], [0.7, 33, 33]] TEST_CASES_2D_2 = [ - {"spatial_dims": 2, "prob_threshold": 0.5, "pixel_dis": 5}, - {"resolution_level": 1}, + {"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[33, 33] = 0.7 +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, 266, 266]] +expected_3 = [[0.9, 66, 66], [0.8, 60, 66]] TEST_CASES_2D_3 = [ - {"spatial_dims": 2, "prob_threshold": 0.5, "pixel_dis": 40}, - {"resolution_level": 2}, + {"spatial_dims": 2, "prob_threshold": 0.5, "box_size": (10, 20)}, probs_map_3, expected_3, ] @@ -45,29 +44,27 @@ 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, 266, 266]] +expected_4 = [[0.9, 66, 66]] TEST_CASES_2D_4 = [ - {"spatial_dims": 2, "prob_threshold": 0.8, "pixel_dis": 5}, - {"resolution_level": 2}, + {"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}, {"resolution_level": 2}, probs_map_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}, {"resolution_level": 2}, probs_map_6, []] +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, 266, 266], [0.7, 134, 134]] +expected_7 = [[0.9, 66, 66], [0.7, 33, 33]] TEST_CASES_2D_7 = [ {"spatial_dims": 2, "prob_threshold": 0.5, "sigma": 0.1}, - {"resolution_level": 2}, probs_map_7, expected_7, ] @@ -77,8 +74,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, "pixel_dis": 5}, - {"resolution_level": 0}, + {"spatial_dims": 3, "prob_threshold": 0.5, "box_size": (10, 10, 10)}, probs_map_3d, expected_3d, ] @@ -97,9 +93,9 @@ class TestPixelNMS(unittest.TestCase): TEST_CASES_3D, ] ) - def test_output(self, class_args, call_args, probs_map, expected): - nms = PixelNMS(**class_args) - output = nms(probs_map, **call_args) + def test_output(self, class_args, probs_map, expected): + nms = ProbNMS(**class_args) + output = nms(probs_map) np.testing.assert_allclose(output, expected) From 4a4d338c157a087f17f26ae87e20d80e658512c6 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Fri, 26 Mar 2021 14:51:28 +0800 Subject: [PATCH 5/7] Add pathology nms Signed-off-by: Yiheng Wang --- monai/apps/pathology/__init__.py | 1 + monai/apps/pathology/utils.py | 64 ++++++++++++++++++++++++++++++++ tests/test_pathology_prob_nms.py | 57 ++++++++++++++++++++++++++++ tests/test_prob_nms.py | 2 +- 4 files changed, 123 insertions(+), 1 deletion(-) create mode 100644 monai/apps/pathology/utils.py create mode 100644 tests/test_pathology_prob_nms.py 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..65eb7206f4 --- /dev/null +++ b/monai/apps/pathology/utils.py @@ -0,0 +1,64 @@ +# 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. + """ + 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 + resolution = pow(2, resolution_level) + + 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) + coord_wsi = ((max_idx_arr + 0.5) * resolution).astype(int) + outputs.append([prob_max] + list(coord_wsi)) + + 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 index ffccf62e2a..fb88d9cfb4 100644 --- a/tests/test_prob_nms.py +++ b/tests/test_prob_nms.py @@ -80,7 +80,7 @@ ] -class TestPixelNMS(unittest.TestCase): +class TestProbNMS(unittest.TestCase): @parameterized.expand( [ TEST_CASES_2D_1, From e1860f4f4b384b36c2285baa5ca369ec44729647 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Fri, 26 Mar 2021 14:57:57 +0800 Subject: [PATCH 6/7] Update docs Signed-off-by: Yiheng Wang --- docs/source/apps.rst | 4 ++++ docs/source/utils.rst | 3 ++- 2 files changed, 6 insertions(+), 1 deletion(-) 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 428160d54d..071d9ecefd 100644 --- a/docs/source/utils.rst +++ b/docs/source/utils.rst @@ -29,7 +29,8 @@ Misc Prob NMS -------- -.. autoclass:: monai.utils.ProbNMS +.. automodule:: monai.utils.prob_nms +.. autoclass:: ProbNMS :members: Profiling From 4306cde5bfca69be0186a6a4952ac4cd1f0ec9f4 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Fri, 26 Mar 2021 15:44:49 +0800 Subject: [PATCH 7/7] Update pathology nms Signed-off-by: Yiheng Wang --- monai/apps/pathology/utils.py | 35 +++++++---------------------------- 1 file changed, 7 insertions(+), 28 deletions(-) diff --git a/monai/apps/pathology/utils.py b/monai/apps/pathology/utils.py index 65eb7206f4..b0803526fd 100644 --- a/monai/apps/pathology/utils.py +++ b/monai/apps/pathology/utils.py @@ -30,35 +30,14 @@ def __call__( ): """ probs_map: the input probabilities map, it must have shape (H[, W, ...]). - resolution_level: the level at which the probabilities map is made. + resolution_level: the level at which the probabilities map is made. """ - 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 resolution = pow(2, resolution_level) - + org_outputs = ProbNMS.__call__(self, probs_map) 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) - coord_wsi = ((max_idx_arr + 0.5) * resolution).astype(int) - outputs.append([prob_max] + list(coord_wsi)) - - 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 - + 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