Add pythonic slicing support to torchio.Image

[marius-sm]

🚀 Feature

Allow torchio.Image to be sliced like numpy arrays

image = torchio.Image(...)
print(image.spatial_shape) # -> (100, 100, 100)

crop = image[10:20, :, ::2] # returns a new ScalarImage where the affine matrix and data have been modified accordingly
print(crop.spatial_shape) # -> (10, 100, 50)

# image['data'] or image['affine'] would still work

If the image is not yet loaded when the slicing is applied, it would modify the affine and when the image is finally loaded, it would load only the given slice, making use of NiBabel's slicing https://nipy.org/nibabel/nibabel_images.html#image-slicing.

Motivation

• More pythonic cropping.
• Easier subsampling and flipping.
• Much faster patch loading for patch-based pipelines.

Additional context

The __getitem__ method of torchio.Image could look like this:

def __getitem__(self, item):
    if isinstance(item, str):
        if item in (DATA, AFFINE):
            if item not in self:
                self.load()
        return super().__getitem__(item)
    else:
        check_validity(item) # do some checks
        if not isinstance(item, tuple):
            item = (item,)
        self.affine = modify_affine(self.affine, item) # to be implemented.
        if not self._loaded:
            self._slice = modify_slice(self._slice, item) # to be implemented. self._slice will be used when the
                                                          # image is finally loaded into memory. It would be
                                                          # initialized as an empty tuple.
        else:
            self.set_data(self.data[(slice(None), *item)]) # add the channel dim
        return self

Whether to modify the image inplace or to return a new image must be discussed.

[fepegar]

Hi, @marius-sm. Apologies, I missed this issue! It sounds like a good idea. Would you like to give it a try?

If the image has already been loaded, using Crop is probably the easiest way to implement this. We could start there.

[fepegar]

I've implemented this in #1170.