Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 56 additions & 4 deletions monai/apps/nuclick/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
"""
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
"""
Expand Down Expand Up @@ -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):
Comment thread
coderabbitai[bot] marked this conversation as resolved.
"""
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):
Expand Down
103 changes: 103 additions & 0 deletions tests/apps/nuclick/test_nuclick_transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]],
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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},
Expand All @@ -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]
Expand Down Expand Up @@ -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):

Expand All @@ -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):

Expand Down