diff --git a/monai/transforms/croppad/array.py b/monai/transforms/croppad/array.py index fc913fa767..5d21e62951 100644 --- a/monai/transforms/croppad/array.py +++ b/monai/transforms/croppad/array.py @@ -827,6 +827,7 @@ def __init__( k_divisible: Sequence[int] | int = 1, mode: str = PytorchPadMode.CONSTANT, lazy: bool = False, + keep_largest_component: bool = False, **pad_kwargs, ) -> None: """ @@ -849,6 +850,9 @@ def __init__( See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html lazy: a flag to indicate whether this transform should execute lazily or not. Defaults to False. + keep_largest_component: if `True`, keep only the largest connected component of the foreground mask + before computing the bounding box, dropping smaller disconnected foreground regions (for example, + isolated text/marker annotations next to the anatomy of interest). Default to `False`. pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. note that `np.pad` treats channel dimension as the first dimension. @@ -860,6 +864,7 @@ def __init__( self.allow_smaller = allow_smaller self.return_coords = return_coords self.k_divisible = k_divisible + self.keep_largest_component = keep_largest_component self.padder = Pad(mode=mode, lazy=lazy, **pad_kwargs) @Crop.lazy.setter # type: ignore @@ -878,7 +883,7 @@ def compute_bounding_box(self, img: torch.Tensor) -> tuple[np.ndarray, np.ndarra """ box_start, box_end = generate_spatial_bounding_box( - img, self.select_fn, self.channel_indices, self.margin, self.allow_smaller + img, self.select_fn, self.channel_indices, self.margin, self.allow_smaller, self.keep_largest_component ) box_start_, *_ = convert_data_type(box_start, output_type=np.ndarray, dtype=np.int16, wrap_sequence=True) box_end_, *_ = convert_data_type(box_end, output_type=np.ndarray, dtype=np.int16, wrap_sequence=True) diff --git a/monai/transforms/croppad/dictionary.py b/monai/transforms/croppad/dictionary.py index d089cea457..d98fa25bc7 100644 --- a/monai/transforms/croppad/dictionary.py +++ b/monai/transforms/croppad/dictionary.py @@ -861,6 +861,7 @@ def __init__( end_coord_key: str | None = "foreground_end_coord", allow_missing_keys: bool = False, lazy: bool = False, + keep_largest_component: bool = False, **pad_kwargs, ) -> None: """ @@ -889,6 +890,9 @@ def __init__( end_coord_key: key to record the end coordinate of spatial bounding box for foreground. allow_missing_keys: don't raise exception if key is missing. lazy: a flag to indicate whether this transform should execute lazily or not. Defaults to False. + keep_largest_component: if `True`, keep only the largest connected component of the foreground mask + before computing the bounding box, dropping smaller disconnected foreground regions (for example, + isolated text/marker annotations next to the anatomy of interest). Default to `False`. pad_kwargs: other arguments for the `np.pad` or `torch.pad` function. note that `np.pad` treats channel dimension as the first dimension. @@ -901,6 +905,7 @@ def __init__( channel_indices=channel_indices, margin=margin, allow_smaller=allow_smaller, + keep_largest_component=keep_largest_component, k_divisible=k_divisible, lazy=lazy, **pad_kwargs, diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 0b8a65b0fb..7786148080 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -1129,6 +1129,7 @@ def generate_spatial_bounding_box( channel_indices: IndexSelection | None = None, margin: Sequence[int] | int = 0, allow_smaller: bool = False, + keep_largest_component: bool = False, ) -> tuple[list[int], list[int]]: """ Generate the spatial bounding box of foreground in the image with start-end positions (inclusive). @@ -1151,12 +1152,19 @@ def generate_spatial_bounding_box( final box edges. If `True`, the bounding boxes edges are aligned with the input image edges, if `False`, the bounding boxes edges are aligned with the final box edges. Default to `False`. The default value is changed from `True` to `False` in v1.5.0. + keep_largest_component: if `True`, keep only the largest connected component of the `select_fn(img)` mask + before computing the bounding box, dropping smaller disconnected foreground regions (for example, + isolated text/marker annotations next to the anatomy of interest). Default to `False` to preserve + existing behavior, since `select_fn` can legitimately select multiple disjoint structures that should + all remain in the box (e.g. several organs in one volume). """ check_non_lazy_pending_ops(img, name="generate_spatial_bounding_box") spatial_size = img.shape[1:] data = img[list(ensure_tuple(channel_indices))] if channel_indices is not None else img data = select_fn(data).any(0) + if keep_largest_component: + data = get_largest_connected_component_mask(data) ndim = len(data.shape) margin = ensure_tuple_rep(margin, ndim) for m in margin: diff --git a/tests/transforms/test_crop_foreground.py b/tests/transforms/test_crop_foreground.py index c533e46ee4..45b9e05d8d 100644 --- a/tests/transforms/test_crop_foreground.py +++ b/tests/transforms/test_crop_foreground.py @@ -88,6 +88,30 @@ ] ) + # a large 3x3 blob plus a single disconnected pixel (e.g. a scanner marker/text annotation) -- + # keep_largest_component=True drops the isolated pixel before computing the box, so only the + # blob is cropped out instead of a box that stretches to cover both. + TESTS.append( + [ + {"select_fn": lambda x: x > 0, "channel_indices": None, "margin": 0, "keep_largest_component": True}, + p( + [ + [ + [0, 0, 0, 0, 0, 0, 0], + [0, 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, 0], + [0, 0, 0, 0, 0, 0, 0], + ] + ] + ), + p([[[1, 1, 1], [1, 1, 1], [1, 1, 1]]]), + True, + ] + ) + TEST_LAZY_ERROR.append( [ {"select_fn": lambda x: x > 0, "channel_indices": None, "margin": 0, "k_divisible": 10}, diff --git a/tests/transforms/test_crop_foregroundd.py b/tests/transforms/test_crop_foregroundd.py index 83d7a8e07c..6b758ce916 100644 --- a/tests/transforms/test_crop_foregroundd.py +++ b/tests/transforms/test_crop_foregroundd.py @@ -155,6 +155,41 @@ False, ] ) + # a large 3x3 blob plus a single disconnected pixel (e.g. a scanner marker/text annotation) -- + # keep_largest_component=True drops the isolated pixel before computing the box, so only the + # blob is cropped out instead of a box that stretches to cover both. Mirrors the equivalent + # case in test_crop_foreground.py, exercised here through the dict transform's forwarding. + TESTS.append( + [ + { + "keys": ["img"], + "source_key": "img", + "select_fn": lambda x: x > 0, + "channel_indices": None, + "margin": 0, + "keep_largest_component": True, + }, + { + "img": p( + np.array( + [ + [ + [0, 0, 0, 0, 0, 0, 0], + [0, 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, 0], + [0, 0, 0, 0, 0, 0, 0], + ] + ] + ) + ) + }, + p(np.array([[[1, 1, 1], [1, 1, 1], [1, 1, 1]]])), + True, + ] + ) class TestCropForegroundd(unittest.TestCase): diff --git a/tests/transforms/test_generate_spatial_bounding_box.py b/tests/transforms/test_generate_spatial_bounding_box.py index 1f63b2c1ef..ad8a2acc15 100644 --- a/tests/transforms/test_generate_spatial_bounding_box.py +++ b/tests/transforms/test_generate_spatial_bounding_box.py @@ -101,6 +101,48 @@ ([0, 0], [5, 5]), ] ) + # a large 3x3 blob plus a single disconnected pixel (e.g. a scanner marker/text annotation + # away from the anatomy of interest) -- without keep_largest_component the box stretches to + # cover both; with it, only the largest connected component (the blob) is kept. + _two_component_img = p( + np.array( + [ + [ + [0, 0, 0, 0, 0, 0, 0], + [0, 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, 0], + [0, 0, 0, 0, 0, 0, 0], + ] + ] + ) + ) + TESTS.append( + [ + { + "img": _two_component_img, + "select_fn": lambda x: x > 0, + "channel_indices": None, + "margin": 0, + "keep_largest_component": False, + }, + ([1, 1], [6, 6]), + ] + ) + TESTS.append( + [ + { + "img": _two_component_img, + "select_fn": lambda x: x > 0, + "channel_indices": None, + "margin": 0, + "keep_largest_component": True, + }, + ([1, 1], [4, 4]), + ] + ) class TestGenerateSpatialBoundingBox(unittest.TestCase):