Add Resize transform

TorchIO-project · fepegar · Sep 2, 2021 · Sep 1, 2021 · Sep 1, 2021 · Sep 1, 2021
commit 761aac14137d49b81e3546e60c0e6d9306b2070b
diff --git a/docs/source/transforms/preprocessing.rst b/docs/source/transforms/preprocessing.rst
@@ -86,6 +86,13 @@ Spatial
     :show-inheritance:
 
 
+:class:`Resize`
+~~~~~~~~~~~~~~~
+
+.. autoclass:: Resize
+    :show-inheritance:
+
+
 :class:`EnsureShapeMultiple`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

diff --git a/tests/transforms/preprocessing/test_resize.py b/tests/transforms/preprocessing/test_resize.py
@@ -0,0 +1,19 @@
+import torchio as tio
+from ...utils import TorchioTestCase
+
+
+class TestResize(TorchioTestCase):
+    """Tests for `Resize`."""
+    def test_one_dim(self):
+        target_shape = 5
+        transform = tio.Resize(target_shape)
+        transformed = transform(self.sample_subject)
+        for image in transformed.get_images(intensity_only=False):
+            self.assertEqual(image.spatial_shape, 3 * (target_shape,))
+
+    def test_all_dims(self):
+        target_shape = 11, 6, 7
+        transform = tio.Resize(target_shape)
+        transformed = transform(self.sample_subject)
+        for image in transformed.get_images(intensity_only=False):
+            self.assertEqual(image.spatial_shape, target_shape)
diff --git a/torchio/transforms/__init__.py b/torchio/transforms/__init__.py
@@ -28,6 +28,7 @@
 # Preprocessing
 from .preprocessing import Pad
 from .preprocessing import Crop
+from .preprocessing import Resize
 from .preprocessing import Resample
 from .preprocessing import CropOrPad
 from .preprocessing import CopyAffine
@@ -82,6 +83,7 @@
     'LabelsToImage',
     'Pad',
     'Crop',
+    'Resize',
     'Resample',
     'ToCanonical',
     'ZNormalization',

diff --git a/torchio/transforms/preprocessing/__init__.py b/torchio/transforms/preprocessing/__init__.py
@@ -1,5 +1,6 @@
 from .spatial.pad import Pad
 from .spatial.crop import Crop
+from .spatial.resize import Resize
 from .spatial.resample import Resample
 from .spatial.crop_or_pad import CropOrPad
 from .spatial.to_canonical import ToCanonical
@@ -23,6 +24,7 @@
 __all__ = [
     'Pad',
     'Crop',
+    'Resize',
     'Resample',
     'ToCanonical',
     'CropOrPad',

diff --git a/torchio/transforms/preprocessing/spatial/crop_or_pad.py b/torchio/transforms/preprocessing/spatial/crop_or_pad.py
@@ -11,7 +11,7 @@
 
 
 class CropOrPad(BoundsTransform):
-    """Crop and/or pad an image to a target shape.
+    """Modify the field of view by cropping or padding to match a target shape.
 
     This transform modifies the affine matrix associated to the volume so that
     physical positions of the voxels are maintained.

diff --git a/torchio/transforms/preprocessing/spatial/resize.py b/torchio/transforms/preprocessing/spatial/resize.py
@@ -0,0 +1,58 @@
+import warnings
+
+import numpy as np
+
+from ....utils import to_tuple
+from ....data.subject import Subject
+from ....typing import TypeSpatialShape
+from ... import SpatialTransform
+from .resample import Resample
+from .crop_or_pad import CropOrPad
+
+
+class Resize(SpatialTransform):
+    """Resample images so the output shape matches the given target shape.
+
+    The field of view remains the same.
+
+    Args:
+        target_shape: Tuple :math:`(W, H, D)`. If a single value :math:`N` is
+            provided, then :math:`W = H = D = N`.
+        image_interpolation: See :ref:`Interpolation`.
+    """
+    def __init__(
+            self,
+            target_shape: TypeSpatialShape,
+            image_interpolation: str = 'linear',
+            **kwargs
+            ):
+        super().__init__(**kwargs)
+        self.target_shape = np.asarray(to_tuple(target_shape, length=3))
+        self.image_interpolation = self.parse_interpolation(
+            image_interpolation)
+        self.args_names = (
+            'target_shape',
+            'image_interpolation',
+        )
+
+    def apply_transform(self, subject: Subject) -> Subject:
+        shape_in = np.asarray(subject.spatial_shape)
+        shape_out = self.target_shape
+        spacing_in = np.asarray(subject.spacing)
+        spacing_out = shape_in / shape_out * spacing_in
+        resample = Resample(
+            spacing_out,
+            image_interpolation=self.image_interpolation,
+        )
+        resampled = resample(subject)
+        # Sometimes, the output shape is one voxel too large
+        # Probably because Resample uses np.ceil to compute the shape
+        if not resampled.spatial_shape == tuple(shape_out):
+            message = (
+                f'Output shape {resampled.spatial_shape}'
+                f' != target shape {tuple(shape_out)}. Fixing with CropOrPad'
+            )
+            warnings.warn(message)
+            crop_pad = CropOrPad(shape_out)
+            resampled = crop_pad(resampled)
+        return resampled