Add opt-in keep_largest_component to CropForeground/CropForegroundd

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)

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,

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:

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},

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):

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):