[DLMED] enhance CropForegroundd transform (#2808)

Signed-off-by: Nic Ma <[email protected]>

Author: Nic-Ma

monai/transforms/croppad/array.py

@@ -691,7 +691,13 @@ def compute_bounding_box(self, img: np.ndarray):
         box_end_ = box_start_ + spatial_size
         return box_start_, box_end_
 
-    def crop_pad(self, img: np.ndarray, box_start: np.ndarray, box_end: np.ndarray):
+    def crop_pad(
+        self,
+        img: np.ndarray,
+        box_start: np.ndarray,
+        box_end: np.ndarray,
+        mode: Optional[Union[NumpyPadMode, str]] = None,
+    ):
         """
         Crop and pad based on the bounding box.
 
@@ -700,15 +706,15 @@ def crop_pad(self, img: np.ndarray, box_start: np.ndarray, box_end: np.ndarray):
         pad_to_start = np.maximum(-box_start, 0)
         pad_to_end = np.maximum(box_end - np.asarray(img.shape[1:]), 0)
         pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist())))
-        return BorderPad(spatial_border=pad, mode=self.mode, **self.np_kwargs)(cropped)
+        return BorderPad(spatial_border=pad, mode=mode or self.mode, **self.np_kwargs)(cropped)
 
-    def __call__(self, img: np.ndarray):
+    def __call__(self, img: np.ndarray, mode: Optional[Union[NumpyPadMode, str]] = None):
         """
         Apply the transform to `img`, assuming `img` is channel-first and
         slicing doesn't change the channel dim.
         """
         box_start, box_end = self.compute_bounding_box(img)
-        cropped = self.crop_pad(img, box_start, box_end)
+        cropped = self.crop_pad(img, box_start, box_end, mode)
 
         if self.return_coords:
             return cropped, box_start, box_end

monai/transforms/croppad/dictionary.py

@@ -797,7 +797,7 @@ def __init__(
         channel_indices: Optional[IndexSelection] = None,
         margin: Union[Sequence[int], int] = 0,
         k_divisible: Union[Sequence[int], int] = 1,
-        mode: Union[NumpyPadMode, str] = NumpyPadMode.CONSTANT,
+        mode: NumpyPadModeSequence = NumpyPadMode.CONSTANT,
         start_coord_key: str = "foreground_start_coord",
         end_coord_key: str = "foreground_end_coord",
         allow_missing_keys: bool = False,
@@ -818,6 +818,7 @@ def __init__(
                 ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
                 one of the listed string values or a user supplied function. Defaults to ``"constant"``.
                 see also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
+                it also can be a sequence of string, each element corresponds to a key in ``keys``.
             start_coord_key: key to record the start coordinate of spatial bounding box for foreground.
             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.
@@ -834,18 +835,18 @@ def __init__(
             channel_indices=channel_indices,
             margin=margin,
             k_divisible=k_divisible,
-            mode=mode,
             **np_kwargs,
         )
+        self.mode = ensure_tuple_rep(mode, len(self.keys))
 
     def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
         d = dict(data)
         box_start, box_end = self.cropper.compute_bounding_box(img=d[self.source_key])
         d[self.start_coord_key] = box_start
         d[self.end_coord_key] = box_end
-        for key in self.key_iterator(d):
+        for key, m in self.key_iterator(d, self.mode):
             self.push_transform(d, key, extra_info={"box_start": box_start, "box_end": box_end})
-            d[key] = self.cropper.crop_pad(d[key], box_start, box_end)
+            d[key] = self.cropper.crop_pad(img=d[key], box_start=box_start, box_end=box_end, mode=m)
         return d
 
     def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:

tests/test_crop_foregroundd.py

@@ -15,6 +15,7 @@
 from parameterized import parameterized
 
 from monai.transforms import CropForegroundd
+from monai.utils import NumpyPadMode
 
 TEST_CASE_1 = [
     {
@@ -59,15 +60,18 @@
 
 TEST_CASE_6 = [
     {
-        "keys": ["img"],
+        "keys": ["img", "seg"],
         "source_key": "img",
         "select_fn": lambda x: x > 0,
         "channel_indices": 0,
         "margin": 0,
         "k_divisible": [4, 6],
-        "mode": "edge",
+        "mode": ["edge", NumpyPadMode.CONSTANT],
+    },
+    {
+        "img": np.array([[[0, 2, 1, 2, 0], [1, 1, 2, 1, 1], [2, 2, 3, 2, 2], [1, 1, 2, 1, 1], [0, 0, 0, 0, 0]]]),
+        "seg": np.array([[[0, 2, 1, 2, 0], [1, 1, 2, 1, 1], [2, 2, 3, 2, 2], [1, 1, 2, 1, 1], [0, 0, 0, 0, 0]]]),
     },
-    {"img": np.array([[[0, 2, 1, 2, 0], [1, 1, 2, 1, 1], [2, 2, 3, 2, 2], [1, 1, 2, 1, 1], [0, 0, 0, 0, 0]]])},
     np.array([[[0, 2, 1, 2, 0, 0], [1, 1, 2, 1, 1, 1], [2, 2, 3, 2, 2, 2], [1, 1, 2, 1, 1, 1]]]),
 ]