Repository: fepegar/torchio

Issue:
Add support for 2D images

Discussion:

Patch:
diff --git a/torchio/data/dataset.py b/torchio/data/dataset.py
index 0bc00f0..6fdf1f2 100644
--- a/torchio/data/dataset.py
+++ b/torchio/data/dataset.py
@@ -207,7 +207,7 @@ class ImagesDataset(Dataset):
             output_paths_dict: Dict[str, TypePath],
             ) -> None:
         for key, output_path in output_paths_dict.items():
-            tensor = sample[key][DATA][0]  # remove channels dim
+            tensor = sample[key][DATA].squeeze()  # assume 2D if (1, 1, H, W)
             affine = sample[key][AFFINE]
             write_image(tensor, affine, output_path)
 
diff --git a/torchio/data/image.py b/torchio/data/image.py
index ccac1dc..f40f4bb 100644
--- a/torchio/data/image.py
+++ b/torchio/data/image.py
@@ -98,8 +98,8 @@ class Image(dict):
     def load(self, check_nans: bool = True) -> Tuple[torch.Tensor, np.ndarray]:
         r"""Load the image from disk.
 
-        The file is expected to be monomodal and 3D. A channels dimension is
-        added to the tensor.
+        The file is expected to be monomodal/grayscale and 2D or 3D.
+        A channels dimension is added to the tensor.
 
         Args:
             check_nans: If ``True``, issues a warning if NaNs are found
@@ -113,6 +113,8 @@ class Image(dict):
         if self.path is None:
             return self.tensor, self.affine
         tensor, affine = read_image(self.path)
+        # https://github.com/pytorch/pytorch/issues/9410#issuecomment-404968513
+        tensor = tensor[(None,) * (3 - tensor.ndim)]  # force to be 3D
         tensor = tensor.unsqueeze(0)  # add channels dimension
         if check_nans and torch.isnan(tensor).any():
             warnings.warn(f'NaNs found in file "{self.path}"')
diff --git a/torchio/data/sampler/sampler.py b/torchio/data/sampler/sampler.py
index edf39a4..8e31ced 100644
--- a/torchio/data/sampler/sampler.py
+++ b/torchio/data/sampler/sampler.py
@@ -97,15 +97,16 @@ def get_random_indices_from_shape(
         shape: Tuple[int, int, int],
         patch_size: Tuple[int, int, int],
         ) -> Tuple[np.ndarray, np.ndarray]:
-    shape_array = np.array(shape, dtype=np.uint16)
-    patch_size_array = np.array(patch_size, dtype=np.uint16)
+    shape_array = np.array(shape)
+    patch_size_array = np.array(patch_size)
     max_index_ini = shape_array - patch_size_array
     if (max_index_ini < 0).any():
         message = (
-            f'Patch size {patch_size_array} must not be'
-            f' larger than image size {tuple(shape_array)}'
+            f'Patch size {patch_size} must not be'
+            f' larger than image size {shape}'
         )
         raise ValueError(message)
+    max_index_ini = max_index_ini.astype(np.uint16)
     coordinates = []
     for max_coordinate in max_index_ini.tolist():
         if max_coordinate == 0:
diff --git a/torchio/transforms/augmentation/intensity/random_motion.py b/torchio/transforms/augmentation/intensity/random_motion.py
index e44efe9..4bbf3c6 100644
--- a/torchio/transforms/augmentation/intensity/random_motion.py
+++ b/torchio/transforms/augmentation/intensity/random_motion.py
@@ -128,9 +128,9 @@ class RandomMotion(RandomTransform):
             degrees_range, num_transforms)
         translation_params = get_params_array(
             translation_range, num_transforms)
-        if is_2d:
-            degrees_params[:-1] = 0  # rotate around z axis only
-            translation_params[-1] = 0  # translate in xy plane only
+        if is_2d:  # imagine sagittal (1, A, S)
+            degrees_params[:, -2:] = 0  # rotate around R axis only
+            translation_params[:, 0] = 0  # translate in AS plane only
         step = 1 / (num_transforms + 1)
         times = torch.arange(0, 1, step)[1:]
         noise = torch.FloatTensor(num_transforms)
diff --git a/torchio/transforms/augmentation/spatial/random_affine.py b/torchio/transforms/augmentation/spatial/random_affine.py
index 9a99137..a096060 100644
--- a/torchio/transforms/augmentation/spatial/random_affine.py
+++ b/torchio/transforms/augmentation/spatial/random_affine.py
@@ -88,22 +88,26 @@ class RandomAffine(RandomTransform):
         sample.check_consistent_shape()
         scaling_params, rotation_params = self.get_params(
             self.scales, self.degrees, self.isotropic)
-        random_parameters_dict = {
-            'scaling': scaling_params,
-            'rotation': rotation_params,
-        }
         for image_dict in sample.get_images(intensity_only=False):
             if image_dict[TYPE] == LABEL:
                 interpolation = Interpolation.NEAREST
             else:
                 interpolation = self.interpolation
+            is_2d = image_dict[DATA].shape[-3] == 1
+            if is_2d:
+                scaling_params[0] = 1
+                rotation_params[-2:] = 0
             image_dict[DATA] = self.apply_affine_transform(
                 image_dict[DATA],
                 image_dict[AFFINE],
-                scaling_params,
-                rotation_params,
+                scaling_params.tolist(),
+                rotation_params.tolist(),
                 interpolation,
             )
+        random_parameters_dict = {
+            'scaling': scaling_params,
+            'rotation': rotation_params,
+        }
         sample.add_transform(self, random_parameters_dict)
         return sample
 
@@ -112,12 +116,12 @@ class RandomAffine(RandomTransform):
             scales: Tuple[float, float],
             degrees: Tuple[float, float],
             isotropic: bool,
-            ) -> Tuple[List[float], List[float]]:
+            ) -> Tuple[np.ndarray, np.ndarray]:
         scaling_params = torch.FloatTensor(3).uniform_(*scales)
         if isotropic:
             scaling_params.fill_(scaling_params[0])
         rotation_params = torch.FloatTensor(3).uniform_(*degrees)
-        return scaling_params.tolist(), rotation_params.tolist()
+        return scaling_params.numpy(), rotation_params.numpy()
 
     @staticmethod
     def get_scaling_transform(
diff --git a/torchio/transforms/augmentation/spatial/random_elastic_deformation.py b/torchio/transforms/augmentation/spatial/random_elastic_deformation.py
index f865c6b..ac5d054 100644
--- a/torchio/transforms/augmentation/spatial/random_elastic_deformation.py
+++ b/torchio/transforms/augmentation/spatial/random_elastic_deformation.py
@@ -218,18 +218,21 @@ class RandomElasticDeformation(RandomTransform):
             self.max_displacement,
             self.num_locked_borders,
         )
-        random_parameters_dict = {'coarse_grid': bspline_params}
         for image_dict in sample.get_images(intensity_only=False):
             if image_dict[TYPE] == LABEL:
                 interpolation = Interpolation.NEAREST
             else:
                 interpolation = self.interpolation
+            is_2d = image_dict[DATA].shape[-3] == 1
+            if is_2d:
+                bspline_params[..., -3] = 0  # no displacement in LR axis
             image_dict[DATA] = self.apply_bspline_transform(
                 image_dict[DATA],
                 image_dict[AFFINE],
                 bspline_params,
                 interpolation,
             )
+        random_parameters_dict = {'coarse_grid': bspline_params}
         sample.add_transform(self, random_parameters_dict)
         return sample
 
diff --git a/torchio/transforms/augmentation/spatial/random_flip.py b/torchio/transforms/augmentation/spatial/random_flip.py
index 4db4a2c..b80f950 100644
--- a/torchio/transforms/augmentation/spatial/random_flip.py
+++ b/torchio/transforms/augmentation/spatial/random_flip.py
@@ -15,6 +15,8 @@ class RandomFlip(RandomTransform):
             computed on a per-axis basis.
         p: Probability that this transform will be applied.
         seed: See :py:class:`~torchio.transforms.augmentation.RandomTransform`.
+
+    .. note:: If the input image is 2D, all axes should be in ``(0, 1)``.
     """
 
     def __init__(
@@ -33,16 +35,30 @@ class RandomFlip(RandomTransform):
     def apply_transform(self, sample: Subject) -> dict:
         axes_to_flip_hot = self.get_params(self.axes, self.flip_probability)
         random_parameters_dict = {'axes': axes_to_flip_hot}
-        for image_dict in sample.get_images(intensity_only=False):
-            tensor = image_dict[DATA]
+        items = sample.get_images_dict(intensity_only=False).items()
+        for image_name, image_dict in items:
+            data = image_dict[DATA]
+            is_2d = data.shape[-3] == 1
             dims = []
             for dim, flip_this in enumerate(axes_to_flip_hot):
                 if not flip_this:
                     continue
                 actual_dim = dim + 1  # images are 4D
+                # If the user is using 2D images and they use (0, 1) for axes,
+                # they probably mean (1, 2). This should make this transform
+                # more user-friendly.
+                if is_2d:
+                    actual_dim += 1
+                if actual_dim > 3:
+                    message = (
+                        f'Image "{image_name}" with shape {data.shape} seems to'
+                        ' be 2D, so all axes must be in (0, 1),'
+                        f' but they are {self.axes}'
+                    )
+                    raise RuntimeError(message)
                 dims.append(actual_dim)
-            tensor = torch.flip(tensor, dims=dims)
-            image_dict[DATA] = tensor
+            data = torch.flip(data, dims=dims)
+            image_dict[DATA] = data
         sample.add_transform(self, random_parameters_dict)
         return sample
 
diff --git a/torchio/utils.py b/torchio/utils.py
index 688f737..1c57310 100644
--- a/torchio/utils.py
+++ b/torchio/utils.py
@@ -17,7 +17,7 @@ from .torchio import (
 )
 
 
-FLIP_XY = np.diag((-1, -1, 1))
+FLIP_XY = np.diag((-1, -1, 1))  # used to switch between LPS and RAS
 
 
 def to_tuple(
@@ -124,10 +124,8 @@ def apply_transform_to_file(
 
 
 def guess_type(string: str) -> Any:
-    """
-    Adapted from
-    https://www.reddit.com/r/learnpython/comments/4599hl/module_to_guess_type_from_a_string/czw3f5s
-    """
+    # Adapted from
+    # https://www.reddit.com/r/learnpython/comments/4599hl/module_to_guess_type_from_a_string/czw3f5s
     string = string.replace(' ', '')
     try:
         value = ast.literal_eval(string)
@@ -151,6 +149,7 @@ def guess_type(string: str) -> Any:
 def get_rotation_and_spacing_from_affine(
         affine: np.ndarray,
         ) -> Tuple[np.ndarray, np.ndarray]:
+    # From https://github.com/nipy/nibabel/blob/master/nibabel/orientations.py
     rotation_zoom = affine[:3, :3]
     spacing = np.sqrt(np.sum(rotation_zoom * rotation_zoom, axis=0))
     rotation = rotation_zoom / spacing
@@ -162,21 +161,32 @@ def nib_to_sitk(data: TypeData, affine: TypeData) -> sitk.Image:
     affine = affine.numpy() if isinstance(affine, torch.Tensor) else affine
     origin = np.dot(FLIP_XY, affine[:3, 3]).astype(np.float64)
     rotation, spacing = get_rotation_and_spacing_from_affine(affine)
-    direction = np.dot(FLIP_XY, rotation).flatten()
+    direction = np.dot(FLIP_XY, rotation)
     image = sitk.GetImageFromArray(array.transpose())
+    if array.ndim == 2:  # ignore first dimension if 2D (1, 1, H, W)
+        direction = direction[1:3, 1:3]
     image.SetOrigin(origin)
     image.SetSpacing(spacing)
-    image.SetDirection(direction)
+    image.SetDirection(direction.flatten())
     return image
 
 
 def sitk_to_nib(image: sitk.Image) -> Tuple[np.ndarray, np.ndarray]:
     data = sitk.GetArrayFromImage(image).transpose()
     spacing = np.array(image.GetSpacing())
-    rotation = np.array(image.GetDirection()).reshape(3, 3)
+    direction = np.array(image.GetDirection())
+    origin = image.GetOrigin()
+    if len(direction) == 9:
+        rotation = direction.reshape(3, 3)
+    elif len(direction) == 4:  # ignore first dimension if 2D (1, 1, H, W)
+        rotation_2d = direction.reshape(2, 2)
+        rotation = np.eye(3)
+        rotation[1:3, 1:3] = rotation_2d
+        spacing = 1, *spacing
+        origin = 0, *origin
     rotation = np.dot(FLIP_XY, rotation)
     rotation_zoom = rotation * spacing
-    translation = np.dot(FLIP_XY, image.GetOrigin())
+    translation = np.dot(FLIP_XY, origin)
     affine = np.eye(4)
     affine[:3, :3] = rotation_zoom
     affine[:3, 3] = translation


Test patch:
diff --git a/tests/transforms/augmentation/test_random_flip.py b/tests/transforms/augmentation/test_random_flip.py
new file mode 100644
index 0000000..ba12745
--- /dev/null
+++ b/tests/transforms/augmentation/test_random_flip.py
@@ -0,0 +1,17 @@
+import torchio
+from ...utils import TorchioTestCase
+
+
+class TestRandomFlip(TorchioTestCase):
+    """Tests for `RandomFlip`."""
+    def test_2d(self):
+        sample = self.make_2d(self.sample)
+        transform = torchio.transforms.RandomFlip(
+            axes=(0, 1), flip_probability=1)
+        transform(sample)
+
+    def test_wrong_axes(self):
+        sample = self.make_2d(self.sample)
+        transform = torchio.transforms.RandomFlip(axes=2, flip_probability=1)
+        with self.assertRaises(RuntimeError):
+            transform(sample)
diff --git a/tests/transforms/test_transforms.py b/tests/transforms/test_transforms.py
index 43b211a..dc71bf6 100644
--- a/tests/transforms/test_transforms.py
+++ b/tests/transforms/test_transforms.py
@@ -9,22 +9,28 @@ from ..utils import TorchioTestCase
 class TestTransforms(TorchioTestCase):
     """Tests for all transforms."""
 
-    def get_transform(self, channels):
+    def get_transform(self, channels, is_3d=True):
         landmarks_dict = {
             channel: np.linspace(0, 100, 13) for channel in channels
         }
-        elastic = torchio.RandomElasticDeformation(max_displacement=1)
+        disp = 1 if is_3d else (0.01, 1, 1)
+        elastic = torchio.RandomElasticDeformation(max_displacement=disp)
+        cp_args = (9, 21, 30) if is_3d else (1, 21, 30)
+        flip_axes = (0, 1, 2) if is_3d else (0, 1)
+        swap_patch = (2, 3, 4) if is_3d else (1, 3, 4)
+        pad_args = (1, 2, 3, 0, 5, 6) if is_3d else (0, 0, 3, 0, 5, 6)
+        crop_args = (3, 2, 8, 0, 1, 4) if is_3d else (0, 0, 8, 0, 1, 4)
         transforms = (
-            torchio.CropOrPad((9, 21, 30)),
+            torchio.CropOrPad(cp_args),
             torchio.ToCanonical(),
             torchio.Resample((1, 1.1, 1.25)),
-            torchio.RandomFlip(axes=(0, 1, 2), flip_probability=1),
+            torchio.RandomFlip(axes=flip_axes, flip_probability=1),
             torchio.RandomMotion(),
             torchio.RandomGhosting(axes=(0, 1, 2)),
             torchio.RandomSpike(),
             torchio.RandomNoise(),
             torchio.RandomBlur(),
-            torchio.RandomSwap(patch_size=2, num_iterations=5),
+            torchio.RandomSwap(patch_size=swap_patch, num_iterations=5),
             torchio.Lambda(lambda x: 2 * x, types_to_apply=torchio.INTENSITY),
             torchio.RandomBiasField(),
             torchio.RescaleIntensity((0, 1)),
@@ -32,17 +38,25 @@ class TestTransforms(TorchioTestCase):
             torchio.HistogramStandardization(landmarks_dict=landmarks_dict),
             elastic,
             torchio.RandomAffine(),
-            torchio.OneOf({torchio.RandomAffine(): 3, elastic: 1}),
-            torchio.Pad((1, 2, 3, 0, 5, 6), padding_mode='constant', fill=3),
-            torchio.Crop((3, 2, 8, 0, 1, 4)),
+            torchio.OneOf({
+                torchio.RandomAffine(): 3,
+                elastic: 1,
+            }),
+            torchio.Pad(pad_args, padding_mode='constant', fill=3),
+            torchio.Crop(crop_args),
         )
         return torchio.Compose(transforms)
 
-    def test_transforms_sample(self):
-        transform = self.get_transform(channels=('t1', 't2'))
-        transform(self.sample)
-
     def test_transforms_tensor(self):
         tensor = torch.rand(2, 4, 5, 8)
         transform = self.get_transform(channels=('channel_0', 'channel_1'))
         transform(tensor)
+
+    def test_transforms_sample_3d(self):
+        transform = self.get_transform(channels=('t1', 't2'), is_3d=True)
+        transform(self.sample)
+
+    def test_transforms_sample_2d(self):
+        transform = self.get_transform(channels=('t1', 't2'), is_3d=False)
+        sample = self.make_2d(self.sample)
+        transform(sample)
diff --git a/tests/utils.py b/tests/utils.py
index 8256cb0..3afcf0f 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,3 +1,4 @@
+import copy
 import shutil
 import random
 import tempfile
@@ -6,7 +7,7 @@ from pathlib import Path
 import numpy as np
 import nibabel as nib
 from torchio.datasets import IXITiny
-from torchio import INTENSITY, LABEL, Image, ImagesDataset, Subject
+from torchio import INTENSITY, LABEL, DATA, Image, ImagesDataset, Subject
 
 
 class TorchioTestCase(unittest.TestCase):
@@ -53,6 +54,12 @@ class TorchioTestCase(unittest.TestCase):
         self.dataset = ImagesDataset(self.subjects_list)
         self.sample = self.dataset[-1]
 
+    def make_2d(self, sample):
+        sample = copy.deepcopy(sample)
+        for image in sample.get_images(intensity_only=False):
+            image[DATA] = image[DATA][:, 0:1, ...]
+        return sample
+
     def get_inconsistent_sample(self):
         """Return a sample containing images of different shape."""
         subject = Subject(