Add support to slice images and subjects (#1170)

* Add support to slice images and subjects

* Update docs

* Address issues

* Ignore mypy errors

Author: fepegar

docs/source/data/image.rst

@@ -22,6 +22,10 @@ instantiation time.
 Instead, the data is only loaded when needed for an operation
 (e.g., if a transform is applied to the image).
 
+Images can be sliced using the standard NumPy / PyTorch slicing syntax.
+This operation updates the coordinates origin in the affine matrix
+correspondingly.
+
 The figure below shows two instances of :class:`Image`.
 The instance of :class:`ScalarImage` contains a 4D tensor representing a
 diffusion MRI, which contains four 3D volumes (one per gradient direction),

docs/source/data/subject.rst

@@ -8,6 +8,11 @@ The :class:`Subject` is a data structure used to store
 images associated with a subject and any other metadata necessary for
 processing.
 
+Subject objects can be sliced using the standard NumPy / PyTorch slicing
+syntax, returning a new subject with sliced images.
+This is only possible if all images in the subject have the same spatial
+shape.
+
 All transforms applied to a :class:`Subject` are saved
 in its :attr:`history` attribute (see :ref:`Reproducibility`).
 

src/torchio/data/image.py

@@ -30,9 +30,11 @@
 from ..typing import TypeDirection3D
 from ..typing import TypePath
 from ..typing import TypeQuartetInt
+from ..typing import TypeSlice
 from ..typing import TypeTripletFloat
 from ..typing import TypeTripletInt
 from ..utils import get_stem
+from ..utils import to_tuple
 from ..utils import guess_external_viewer
 from ..utils import is_iterable
 from .io import check_uint_to_int
@@ -201,6 +203,9 @@ def __repr__(self):
         return string
 
     def __getitem__(self, item):
+        if isinstance(item, (slice, int, tuple)):
+            return self._crop_from_slices(item)
+
         if item in (DATA, AFFINE):
             if item not in self:
                 self.load()
@@ -784,6 +789,44 @@ def show(self, viewer_path: Optional[TypePath] = None) -> None:
             image_viewer.SetApplication(str(viewer_path))
             image_viewer.Execute(sitk_image)
 
+    def _crop_from_slices(
+        self,
+        slices: Union[TypeSlice, Tuple[TypeSlice, ...]],
+    ) -> 'Image':
+        from ..transforms import Crop
+
+        slices_tuple = to_tuple(slices)  # type: ignore[assignment]
+        cropping: List[int] = []
+        for dim, slice_ in enumerate(slices_tuple):
+            if isinstance(slice_, slice):
+                pass
+            elif slice_ is Ellipsis:
+                message = 'Ellipsis slicing is not supported yet'
+                raise NotImplementedError(message)
+            elif isinstance(slice_, int):
+                slice_ = slice(slice_, slice_ + 1)  # type: ignore[assignment]
+            else:
+                message = f'Slice type not understood: "{type(slice_)}"'
+                raise TypeError(message)
+            shape_dim = self.spatial_shape[dim]
+            assert isinstance(slice_, slice)
+            start, stop, step = slice_.indices(shape_dim)
+            if step != 1:
+                message = (
+                    'Slicing with steps different from 1 is not supported yet.'
+                    ' Use the Crop transform instead'
+                )
+                raise ValueError(message)
+            crop_ini = start
+            crop_fin = shape_dim - stop
+            cropping.extend([crop_ini, crop_fin])
+        while dim < 2:
+            cropping.extend([0, 0])
+            dim += 1
+        w_ini, w_fin, h_ini, h_fin, d_ini, d_fin = cropping
+        cropping_arg = w_ini, w_fin, h_ini, h_fin, d_ini, d_fin  # making mypy happy
+        return Crop(cropping_arg)(self)  # type: ignore[return-value]
+
 
 class ScalarImage(Image):
     """Image whose pixel values represent scalars.

src/torchio/data/subject.py

@@ -77,6 +77,23 @@ def __copy__(self):
     def __len__(self):
         return len(self.get_images(intensity_only=False))
 
+    def __getitem__(self, item):
+        if isinstance(item, (slice, int, tuple)):
+            try:
+                self.check_consistent_spatial_shape()
+            except RuntimeError as e:
+                message = (
+                    'To use indexing, all images in the subject must have the'
+                    ' same spatial shape'
+                )
+                raise RuntimeError(message) from e
+            copied = copy.deepcopy(self)
+            for image_name, image in copied.items():
+                copied[image_name] = image[item]
+            return copied
+        else:
+            return super().__getitem__(item)
+
     @staticmethod
     def _parse_images(images: List[Image]) -> None:
         # Check that it's not empty

src/torchio/transforms/preprocessing/spatial/resample.py

@@ -158,14 +158,17 @@ def apply_transform(self, subject: Subject) -> Subject:
             self.check_affine_key_presence(self.pre_affine_name, subject)
 
         for image in self.get_images(subject):
-            # Do not resample the reference image if it is in the subject
+            # If the current image is the reference, don't resample it
             if self.target is image:
                 continue
+
+            # If the target is not a string, or is not an image in the subject,
+            # do nothing
             try:
                 target_image = subject[self.target]
                 if target_image is image:
                     continue
-            except (KeyError, TypeError):
+            except (KeyError, TypeError, RuntimeError):
                 pass
 
             # Choose interpolation

src/torchio/typing.py

@@ -15,6 +15,7 @@
 TypeKeys = Optional[Sequence[str]]
 TypeData = Union[torch.Tensor, np.ndarray]
 TypeDataAffine = Tuple[torch.Tensor, np.ndarray]
+TypeSlice = Union[int, slice]
 
 TypeDoubletInt = Tuple[int, int]
 TypeTripletInt = Tuple[int, int, int]

tests/data/test_image.py

@@ -283,3 +283,27 @@ def test_copy_no_data(self):
         new_image = copy.copy(my_image)
         assert my_image._loaded
         assert new_image._loaded
+
+    def test_slicing(self):
+        path = self.get_image_path('im_slicing')
+        image = tio.ScalarImage(path)
+
+        assert image.shape == (1, 10, 20, 30)
+
+        cropped = image[0]
+        assert cropped.shape == (1, 1, 20, 30)
+
+        cropped = image[:, 2:-3]
+        assert cropped.shape == (1, 10, 15, 30)
+
+        cropped = image[-5:, 5:]
+        assert cropped.shape == (1, 5, 15, 30)
+
+        with pytest.raises(NotImplementedError):
+            image[..., 5]
+
+        with pytest.raises(ValueError):
+            image[0:8:-1]
+
+        with pytest.raises(ValueError):
+            image[3::-1]