diff --git a/monai/apps/nuclick/transforms.py b/monai/apps/nuclick/transforms.py index 6b6542308a7..a0b29f99675 100644 --- a/monai/apps/nuclick/transforms.py +++ b/monai/apps/nuclick/transforms.py @@ -137,7 +137,8 @@ class SplitLabeld(MapTransform): label: key of the label source others: other labels storage key, defaults to ``"others"`` mask_value: the mask_value that will be kept for binarization of the label, defaults to ``"mask_value"`` - min_area: The smallest allowable object size. + min_area: The smallest allowable object size. Connected components of ``others`` + smaller than ``min_area`` pixels are discarded during relabeling. others_value: Value/class for other nuclei; Use this to separate core nuclei vs others. to_binary_mask: Convert mask to binary; Set it false to restore original class values """ @@ -184,6 +185,14 @@ def __call__(self, data): others[others > 0] = 1 if torch.count_nonzero(others): others = measure.label(convert_to_numpy(others)[0], connectivity=1) + # discard ``others`` components smaller than ``min_area`` pixels + # (bincount-based so the semantics stay identical across skimage + # versions: skimage 0.26 changed remove_small_objects' threshold + # from "smaller than" to "smaller than or equal to") + sizes = np.bincount(others.ravel()) + too_small = sizes < self.min_area + too_small[0] = False + others = np.where(too_small[others], 0, others) others = torch.from_numpy(others)[None] label = mask.type(torch.uint8) if isinstance(mask, torch.Tensor) else mask @@ -525,10 +534,17 @@ class PostFilterLabeld(MapTransform): Performs Filtering of Labels on the predicted probability map Args: + nuc_points: key of the click-point maps, used as reconstruction markers when + ``do_reconstruction`` is True. + bounding_boxes: key of the bounding boxes for each predicted instance. + img_height: key of the output image height. + img_width: key of the output image width. thresh: probability threshold for classifying a pixel as a mask min_size: min_size objects that will be removed from the image, refer skimage remove_small_objects min_hole: min_hole that will be removed from the image, refer skimage remove_small_holes - do_reconstruction: Boolean Flag, Perform a morphological reconstruction of an image, refer skimage + do_reconstruction: Boolean Flag, Perform a morphological reconstruction of an image, refer skimage. + When enabled, only the mask components containing click points (from the ``nuc_points`` + key) are kept and regrown; components without a click point are discarded. allow_missing_keys: don't raise exception if key is missing. pred_classes: List of Predicted class for each instance """ @@ -569,15 +585,51 @@ def __call__(self, data): for key in self.keys: label = d[key].astype(np.uint8) - masks = self.post_processing(label, self.thresh, self.min_size, self.min_hole) + masks = self.post_processing( + label, + self.thresh, + self.min_size, + self.min_hole, + do_reconstruction=self.do_reconstruction, + nuc_points=d.get(self.nuc_points) if self.do_reconstruction else None, + ) d[key] = self.gen_instance_map(masks, bounding_boxes, x, y, pred_classes=pred_classes).astype(np.uint8) return d - def post_processing(self, preds, thresh=0.33, min_size=10, min_hole=30): + def post_processing(self, preds, thresh=0.33, min_size=10, min_hole=30, do_reconstruction=False, nuc_points=None): + """ + Convert predicted probability maps into cleaned binary instance masks. + + Each channel is thresholded at ``thresh``, small objects and holes are removed, and, + when enabled, the mask is morphologically reconstructed from the nuclear marker points. + + Args: + preds: predicted probability maps, as an array of shape ``(n_instances, H, W)``. + thresh: threshold used to binarize the predictions. + min_size: minimum area for a connected component to be kept, + passed to ``morphology.remove_small_objects``. + min_hole: maximum area of the holes to be filled, + passed to ``morphology.remove_small_holes``. + do_reconstruction: whether to morphologically reconstruct each mask from the + corresponding nuclear marker points. + nuc_points: optional nuclear points, one entry per instance, used as + reconstruction markers when ``do_reconstruction`` is enabled. + + Returns: + Binary masks with the same shape as ``preds``. + """ masks = preds > thresh for i in range(preds.shape[0]): masks[i] = morphology.remove_small_objects(masks[i], min_size=min_size) masks[i] = morphology.remove_small_holes(masks[i], area_threshold=min_hole) + if do_reconstruction and nuc_points is not None and i < len(nuc_points): + points = convert_to_numpy(nuc_points[i]) + marker = points[0] > 0 if points.ndim == 3 else points > 0 + # intersect with the filtered mask so that the marker always + # satisfies skimage's ``marker <= mask`` requirement + marker = np.logical_and(marker, masks[i]) + if np.any(marker): + masks[i] = morphology.reconstruction(marker, masks[i], footprint=morphology.disk(1)) return masks def gen_instance_map(self, masks, bounding_boxes, x, y, flatten=True, pred_classes=None): diff --git a/tests/apps/nuclick/test_nuclick_transforms.py b/tests/apps/nuclick/test_nuclick_transforms.py index a6e66c36586..20ef5fdfc05 100644 --- a/tests/apps/nuclick/test_nuclick_transforms.py +++ b/tests/apps/nuclick/test_nuclick_transforms.py @@ -65,6 +65,11 @@ LABEL_4 = np.array([[[4, 4, 4, 4], [4, 4, 4, 4], [4, 4, 4, 4], [4, 4, 4, 4]]], dtype=np.uint8) +# mask_value block (class 1, area 9), a 1-pixel object (class 2), a 9-pixel object (class 3) +LABEL_5 = np.array( + [[[1, 1, 1, 0, 2, 0, 3, 3, 3], [1, 1, 1, 0, 0, 0, 3, 3, 3], [1, 1, 1, 0, 0, 0, 3, 3, 3]]], dtype=np.uint8 +) + IL_IMAGE_1 = np.array( [ [[0, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1]], @@ -125,6 +130,72 @@ "img_width": 6, } +# two above-threshold blobs: a 3x3 block (area 9) and a 2x2 block (area 4) +PRED_2 = np.array( + [ + [ + [1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 1, 1, 0], + [0, 0, 0, 0, 1, 1, 0], + [0, 0, 0, 0, 0, 0, 0], + ] + ], + dtype=np.float32, +) +# click point inside the smaller blob +NUC_POINTS_2 = np.array( + [ + [ + [ + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 1, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + ] + ] + ], + dtype=np.float32, +) +# click point on background (outside any mask component) +NUC_POINTS_3 = np.array( + [ + [ + [ + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 1, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + ] + ] + ], + dtype=np.float32, +) +BB_2 = np.array([[0, 0, 7, 7]], dtype=np.uint8) + +DATA_LABEL_FILTER_2 = { + "pred": PRED_2, + "nuc_points": NUC_POINTS_2, + "bounding_boxes": BB_2, + "img_height": 7, + "img_width": 7, +} +DATA_LABEL_FILTER_3 = { + "pred": PRED_2, + "nuc_points": NUC_POINTS_3, + "bounding_boxes": BB_2, + "img_height": 7, + "img_width": 7, +} + # Result Definitions EXTRACT_RESULT_TC1 = np.array([[[0, 0, 0], [0, 0, 0], [0, 0, 1]]], dtype=np.uint8) EXTRACT_RESULT_TC2 = np.array([[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]], dtype=np.uint8) @@ -155,6 +226,19 @@ SPLIT_TEST_CASE_1 = [{"keys": ["label"], "mask_value": "mask_value", "min_area": 1}, DATA_SPLIT_1, SPLIT_RESULT_TC1] SPLIT_TEST_CASE_2 = [{"keys": ["label"], "mask_value": "mask_value", "min_area": 3}, DATA_SPLIT_2, SPLIT_RESULT_TC2] +# the 1-pixel "others" object is dropped with min_area=5 (the surviving object keeps +# its original component label 2), kept with min_area=1 +SPLIT_MIN_AREA_CASE_1 = [ + {"keys": ["label"], "mask_value": "mask_value", "min_area": 5}, + {"label": LABEL_5, "mask_value": 1}, + [0, 2], +] +SPLIT_MIN_AREA_CASE_2 = [ + {"keys": ["label"], "mask_value": "mask_value", "min_area": 1}, + {"label": LABEL_5, "mask_value": 1}, + [0, 1, 2], +] + GUIDANCE_TEST_CASE_1 = [{"image": "image", "label": "label", "others": "others"}, DATA_GUIDANCE_1, [5, 5, 5]] GUIDANCE_TEST_CASE_2 = [ {"image": "image", "label": "label", "others": "others", "gaussian": True, "use_distance": True}, @@ -171,6 +255,13 @@ LABEL_FILTER_TEST_CASE_1 = [{"keys": ["pred"]}, DATA_LABEL_FILTER_1, [6, 6]] +# without reconstruction both blobs survive (9 + 4 = 13 pixels); with reconstruction +# only the blob containing the click point is kept (4 pixels) +LABEL_FILTER_TEST_CASE_2 = [{"keys": ["pred"], "min_size": 3}, DATA_LABEL_FILTER_2, 13] +LABEL_FILTER_TEST_CASE_3 = [{"keys": ["pred"], "min_size": 3, "do_reconstruction": True}, DATA_LABEL_FILTER_2, 4] +# a click point outside every mask component leaves the filtered mask unchanged +LABEL_FILTER_TEST_CASE_4 = [{"keys": ["pred"], "min_size": 3, "do_reconstruction": True}, DATA_LABEL_FILTER_3, 13] + LABEL_GUIDANCE_TEST_CASE_1 = [{"keys": ["image"], "source": "label"}, DATA_GUIDANCE_1, [4, 5, 5]] LABEL_CLASS_TEST_CASE_1 = [{"keys": ["label"], "offset": 2}, DATA_GUIDANCE_1, 3] @@ -214,6 +305,12 @@ def test_correct_results(self, arguments, input_data, expected_result): result = SplitLabeld(**arguments)(input_data) np.testing.assert_equal(result["label"], expected_result) + @parameterized.expand([SPLIT_MIN_AREA_CASE_1, SPLIT_MIN_AREA_CASE_2]) + def test_min_area_filters_others(self, arguments, input_data, expected_values): + """Test that ``others`` objects below ``min_area`` are dropped from the output.""" + result = SplitLabeld(**arguments)(input_data) + np.testing.assert_equal(np.unique(result["others"]), expected_values) + class TestGuidanceSignal(unittest.TestCase): @@ -238,6 +335,12 @@ def test_correct_shape(self, arguments, input_data, expected_shape): result = PostFilterLabeld(**arguments)(input_data) np.testing.assert_equal(result["pred"].shape, expected_shape) + @parameterized.expand([LABEL_FILTER_TEST_CASE_2, LABEL_FILTER_TEST_CASE_3, LABEL_FILTER_TEST_CASE_4]) + def test_do_reconstruction(self, arguments, input_data, expected_count): + """Test that reconstruction keeps only the blob containing the click point.""" + result = PostFilterLabeld(**arguments)(input_data) + np.testing.assert_equal(np.count_nonzero(result["pred"]), expected_count) + class TestAddLabelAsGuidance(unittest.TestCase):