allow slicing with SpatialCrop(d) (#1954)

* allow slicing with SpatialCrop(d)

Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com>

Author: rijobro

monai/apps/deepgrow/transforms.py

@@ -698,8 +698,10 @@ def __call__(self, data):
             cropper = SpatialCrop(roi_center=center, roi_size=spatial_size)
         else:
             cropper = SpatialCrop(roi_start=box_start, roi_end=box_end)
-        box_start, box_end = cropper.roi_start, cropper.roi_end
 
+        # update bounding box in case it was corrected by the SpatialCrop constructor
+        box_start = np.array([s.start for s in cropper.slices])
+        box_end = np.array([s.stop for s in cropper.slices])
         for key in self.key_iterator(d):
             if not np.array_equal(d[key].shape[1:], original_spatial_shape):
                 raise RuntimeError("All the image specified in keys should have same spatial shape")

monai/transforms/croppad/array.py

@@ -214,8 +214,11 @@ class SpatialCrop(Transform):
     """
     General purpose cropper to produce sub-volume region of interest (ROI).
     It can support to crop ND spatial (channel-first) data.
-    Either a spatial center and size must be provided, or alternatively,
-    if center and size are not provided, the start and end coordinates of the ROI must be provided.
+
+    The cropped region can be parameterised in various ways:
+        - a list of slices for each spatial dimension (allows for use of -ve indexing and `None`)
+        - a spatial center and size
+        - the start and end coordinates of the ROI
     """
 
     def __init__(
@@ -224,35 +227,44 @@ def __init__(
         roi_size: Union[Sequence[int], np.ndarray, None] = None,
         roi_start: Union[Sequence[int], np.ndarray, None] = None,
         roi_end: Union[Sequence[int], np.ndarray, None] = None,
+        roi_slices: Optional[Sequence[slice]] = None,
     ) -> None:
         """
         Args:
             roi_center: voxel coordinates for center of the crop ROI.
             roi_size: size of the crop ROI.
             roi_start: voxel coordinates for start of the crop ROI.
             roi_end: voxel coordinates for end of the crop ROI.
+            roi_slices: list of slices for each of the spatial dimensions.
         """
-        if roi_center is not None and roi_size is not None:
-            roi_center = np.asarray(roi_center, dtype=np.int16)
-            roi_size = np.asarray(roi_size, dtype=np.int16)
-            self.roi_start = np.maximum(roi_center - np.floor_divide(roi_size, 2), 0)
-            self.roi_end = np.maximum(self.roi_start + roi_size, self.roi_start)
+        if roi_slices:
+            if not all(s.step is None or s.step == 1 for s in roi_slices):
+                raise ValueError("Only slice steps of 1/None are currently supported")
+            self.slices = list(roi_slices)
         else:
-            if roi_start is None or roi_end is None:
-                raise ValueError("Please specify either roi_center, roi_size or roi_start, roi_end.")
-            self.roi_start = np.maximum(np.asarray(roi_start, dtype=np.int16), 0)
-            self.roi_end = np.maximum(np.asarray(roi_end, dtype=np.int16), self.roi_start)
-        # Allow for 1D by converting back to np.array (since np.maximum will convert to int)
-        self.roi_start = self.roi_start if isinstance(self.roi_start, np.ndarray) else np.array([self.roi_start])
-        self.roi_end = self.roi_end if isinstance(self.roi_end, np.ndarray) else np.array([self.roi_end])
+            if roi_center is not None and roi_size is not None:
+                roi_center = np.asarray(roi_center, dtype=np.int16)
+                roi_size = np.asarray(roi_size, dtype=np.int16)
+                roi_start_np = np.maximum(roi_center - np.floor_divide(roi_size, 2), 0)
+                roi_end_np = np.maximum(roi_start_np + roi_size, roi_start_np)
+            else:
+                if roi_start is None or roi_end is None:
+                    raise ValueError("Please specify either roi_center, roi_size or roi_start, roi_end.")
+                roi_start_np = np.maximum(np.asarray(roi_start, dtype=np.int16), 0)
+                roi_end_np = np.maximum(np.asarray(roi_end, dtype=np.int16), roi_start_np)
+            # Allow for 1D by converting back to np.array (since np.maximum will convert to int)
+            roi_start_np = roi_start_np if isinstance(roi_start_np, np.ndarray) else np.array([roi_start_np])
+            roi_end_np = roi_end_np if isinstance(roi_end_np, np.ndarray) else np.array([roi_end_np])
+            # convert to slices
+            self.slices = [slice(s, e) for s, e in zip(roi_start_np, roi_end_np)]
 
     def __call__(self, img: Union[np.ndarray, torch.Tensor]):
         """
         Apply the transform to `img`, assuming `img` is channel-first and
         slicing doesn't apply to the channel dim.
         """
-        sd = min(self.roi_start.size, self.roi_end.size, len(img.shape[1:]))  # spatial dims
-        slices = [slice(None)] + [slice(s, e) for s, e in zip(self.roi_start[:sd], self.roi_end[:sd])]
+        sd = min(len(self.slices), len(img.shape[1:]))  # spatial dims
+        slices = [slice(None)] + self.slices[:sd]
         return img[tuple(slices)]
 
 

monai/transforms/croppad/dictionary.py

@@ -286,8 +286,13 @@ def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndar
 class SpatialCropd(MapTransform, InvertibleTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.SpatialCrop`.
-    Either a spatial center and size must be provided, or alternatively if center and size
-    are not provided, the start and end coordinates of the ROI must be provided.
+    General purpose cropper to produce sub-volume region of interest (ROI).
+    It can support to crop ND spatial (channel-first) data.
+
+    The cropped region can be parameterised in various ways:
+        - a list of slices for each spatial dimension (allows for use of -ve indexing and `None`)
+        - a spatial center and size
+        - the start and end coordinates of the ROI
     """
 
     def __init__(
@@ -297,6 +302,7 @@ def __init__(
         roi_size: Optional[Sequence[int]] = None,
         roi_start: Optional[Sequence[int]] = None,
         roi_end: Optional[Sequence[int]] = None,
+        roi_slices: Optional[Sequence[slice]] = None,
         allow_missing_keys: bool = False,
     ) -> None:
         """
@@ -307,10 +313,11 @@ def __init__(
             roi_size: size of the crop ROI.
             roi_start: voxel coordinates for start of the crop ROI.
             roi_end: voxel coordinates for end of the crop ROI.
+            roi_slices: list of slices for each of the spatial dimensions.
             allow_missing_keys: don't raise exception if key is missing.
         """
         super().__init__(keys, allow_missing_keys)
-        self.cropper = SpatialCrop(roi_center, roi_size, roi_start, roi_end)
+        self.cropper = SpatialCrop(roi_center, roi_size, roi_start, roi_end, roi_slices)
 
     def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
         d = dict(data)
@@ -325,9 +332,11 @@ def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndar
         for key in self.key_iterator(d):
             transform = self.get_most_recent_transform(d, key)
             # Create inverse transform
-            orig_size = transform[InverseKeys.ORIG_SIZE]
-            pad_to_start = np.array(self.cropper.roi_start)
-            pad_to_end = orig_size - self.cropper.roi_end
+            orig_size = np.array(transform[InverseKeys.ORIG_SIZE])
+            current_size = np.array(d[key].shape[1:])
+            # get required pad to start and end
+            pad_to_start = np.array([s.indices(o)[0] for s, o in zip(self.cropper.slices, orig_size)])
+            pad_to_end = orig_size - current_size - pad_to_start
             # interleave mins and maxes
             pad = list(chain(*zip(pad_to_start.tolist(), pad_to_end.tolist())))
             inverse_transform = BorderPad(pad)

tests/test_inverse.py

@@ -122,7 +122,6 @@
     )
 )
 
-
 TESTS.append(
     (
         "SpatialCropd 2d",
@@ -132,6 +131,15 @@
     )
 )
 
+TESTS.append(
+    (
+        "SpatialCropd 3d",
+        "3D",
+        0,
+        SpatialCropd(KEYS, roi_slices=[slice(s, e) for s, e in zip([None, None, -99], [None, -2, None])]),
+    )
+)
+
 TESTS.append(
     (
         "SpatialCropd 2d",

tests/test_spatial_crop.py

@@ -40,6 +40,15 @@
         (3, 3, 3, 3),
         (3, 0, 3, 3),
     ],
+    [
+        {"roi_slices": [slice(s, e) for s, e in zip([-1, -2, 0], [None, None, 2])]},
+        (3, 3, 3, 3),
+        (3, 1, 2, 2),
+    ],
+]
+
+TEST_ERRORS = [
+    [{"roi_slices": [slice(s, e, 2) for s, e in zip([-1, -2, 0], [None, None, 2])]}],
 ]
 
 
@@ -56,6 +65,11 @@ def test_tensor_shape(self, input_param, input_shape, expected_shape):
         result = SpatialCrop(**input_param)(input_data)
         self.assertTupleEqual(result.shape, expected_shape)
 
+    @parameterized.expand(TEST_ERRORS)
+    def test_error(self, input_param):
+        with self.assertRaises(ValueError):
+            SpatialCrop(**input_param)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_spatial_cropd.py

@@ -16,33 +16,37 @@
 
 from monai.transforms import SpatialCropd
 
-TEST_CASE_1 = [
-    {"keys": ["img"], "roi_center": [1, 1, 1], "roi_size": [2, 2, 2]},
-    {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])},
-    (3, 2, 2, 2),
-]
-
-TEST_CASE_2 = [
-    {"keys": ["img"], "roi_start": [0, 0, 0], "roi_end": [2, 2, 2]},
-    {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])},
-    (3, 2, 2, 2),
-]
-
-TEST_CASE_3 = [
-    {"keys": ["img"], "roi_start": [0, 0], "roi_end": [2, 2]},
-    {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])},
-    (3, 2, 2, 3),
-]
-
-TEST_CASE_4 = [
-    {"keys": ["img"], "roi_start": [0, 0, 0, 0, 0], "roi_end": [2, 2, 2, 2, 2]},
-    {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])},
-    (3, 2, 2, 2),
+TEST_CASES = [
+    [
+        {"keys": ["img"], "roi_center": [1, 1, 1], "roi_size": [2, 2, 2]},
+        {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])},
+        (3, 2, 2, 2),
+    ],
+    [
+        {"keys": ["img"], "roi_start": [0, 0, 0], "roi_end": [2, 2, 2]},
+        {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])},
+        (3, 2, 2, 2),
+    ],
+    [
+        {"keys": ["img"], "roi_start": [0, 0], "roi_end": [2, 2]},
+        {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])},
+        (3, 2, 2, 3),
+    ],
+    [
+        {"keys": ["img"], "roi_start": [0, 0, 0, 0, 0], "roi_end": [2, 2, 2, 2, 2]},
+        {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])},
+        (3, 2, 2, 2),
+    ],
+    [
+        {"keys": ["img"], "roi_slices": [slice(s, e) for s, e in zip([-1, -2, 0], [None, None, 2])]},
+        {"img": np.random.randint(0, 2, size=[3, 3, 3, 3])},
+        (3, 1, 2, 2),
+    ],
 ]
 
 
 class TestSpatialCropd(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
+    @parameterized.expand(TEST_CASES)
     def test_shape(self, input_param, input_data, expected_shape):
         result = SpatialCropd(**input_param)(input_data)
         self.assertTupleEqual(result["img"].shape, expected_shape)