Repository: fepegar/torchio

Issue:
ValueError in Subject.spatial_shape
**🐛Bug**
<!-- A clear and concise description of what the bug is -->

An error is raised if images have different number of channels, even if the spatial shape of all images are the same. This happens because it is computed using Subject.shape. It used to work until we added support for 4D images.

**To reproduce**

<!-- Try to provide a minimal working example: https://stackoverflow.com/help/minimal-reproducible-example -->

```python
import torchio as tio
icbm = tio.datasets.ICBM2009CNonlinearSymmetryc()
del icbm['face']  # this one does have a different spatial shape
icbm.spatial_shape
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

```python-traceback
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-2-6b7dc2edb3cc> in <module>
----> 1 icbm.spatial_shape

~/git/torchio/torchio/data/subject.py in spatial_shape(self)
     95         Consistency of shapes across images in the subject is checked first.
     96         """
---> 97         return self.shape[1:]
     98
     99     @property

~/git/torchio/torchio/data/subject.py in shape(self)
     85         Consistency of shapes across images in the subject is checked first.
     86         """
---> 87         self.check_consistent_shape()
     88         image = self.get_images(intensity_only=False)[0]
     89         return image.shape

~/git/torchio/torchio/data/subject.py in check_consistent_shape(self)
    135                 f'\n{pprint.pformat(shapes_dict)}'
    136             )
--> 137             raise ValueError(message)
    138
    139     def check_consistent_orientation(self) -> None:

ValueError: Images in subject have inconsistent shapes:
{'brain': (1, 193, 229, 193),
 'eyes': (1, 193, 229, 193),
 'pd': (1, 193, 229, 193),
 't1': (1, 193, 229, 193),
 't2': (1, 193, 229, 193),
 'tissues': (3, 193, 229, 193)}
```


**Expected behavior**

<!-- A clear and concise description of what you expected to happen. -->

It should check for spatial shapes of images only.


**TorchIO version**

`0.17.26`.


Comments:

Patch:
diff --git a/torchio/data/subject.py b/torchio/data/subject.py
index 0430423..76445e9 100644
--- a/torchio/data/subject.py
+++ b/torchio/data/subject.py
@@ -45,18 +45,15 @@ class Subject(dict):
                     'Only one dictionary as positional argument is allowed')
                 raise ValueError(message)
         super().__init__(**kwargs)
-        self.images = [
-            (k, v) for (k, v) in self.items()
-            if isinstance(v, Image)
-        ]
-        self._parse_images(self.images)
+        self._parse_images(self.get_images(intensity_only=False))
         self.update_attributes()  # this allows me to do e.g. subject.t1
         self.history = []
 
     def __repr__(self):
+        num_images = len(self.get_images(intensity_only=False))
         string = (
             f'{self.__class__.__name__}'
-            f'(Keys: {tuple(self.keys())}; images: {len(self.images)})'
+            f'(Keys: {tuple(self.keys())}; images: {num_images})'
         )
         return string
 
@@ -84,27 +81,46 @@ class Subject(dict):
 
         Consistency of shapes across images in the subject is checked first.
         """
-        self.check_consistent_shape()
-        image = self.get_images(intensity_only=False)[0]
-        return image.shape
+        self.check_consistent_attribute('shape')
+        return self.get_first_image().shape
 
     @property
     def spatial_shape(self):
         """Return spatial shape of first image in subject.
 
-        Consistency of shapes across images in the subject is checked first.
+        Consistency of spatial shapes across images in the subject is checked
+        first.
         """
-        return self.shape[1:]
+        self.check_consistent_spatial_shape()
+        return self.get_first_image().spatial_shape
 
     @property
     def spacing(self):
         """Return spacing of first image in subject.
 
-        Consistency of shapes across images in the subject is checked first.
+        Consistency of spacings across images in the subject is checked first.
         """
-        self.check_consistent_shape()
-        image = self.get_images(intensity_only=False)[0]
-        return image.spacing
+        self.check_consistent_attribute('spacing')
+        return self.get_first_image().spacing
+
+    def check_consistent_attribute(self, attribute: str) -> None:
+        values_dict = {}
+        iterable = self.get_images_dict(intensity_only=False).items()
+        for image_name, image in iterable:
+            values_dict[image_name] = getattr(image, attribute)
+        num_unique_values = len(set(values_dict.values()))
+        if num_unique_values > 1:
+            message = (
+                f'More than one {attribute} found in subject images:'
+                f'\n{pprint.pformat(values_dict)}'
+            )
+            raise RuntimeError(message)
+
+    def check_consistent_shape(self) -> None:
+        self.check_consistent_attribute('shape')
+
+    def check_consistent_spatial_shape(self) -> None:
+        self.check_consistent_attribute('spatial_shape')
 
     def get_images_dict(self, intensity_only=True):
         images = {}
@@ -123,32 +139,6 @@ class Subject(dict):
     def get_first_image(self):
         return self.get_images(intensity_only=False)[0]
 
-    def check_consistent_shape(self) -> None:
-        shapes_dict = {}
-        iterable = self.get_images_dict(intensity_only=False).items()
-        for image_name, image in iterable:
-            shapes_dict[image_name] = image.shape
-        num_unique_shapes = len(set(shapes_dict.values()))
-        if num_unique_shapes > 1:
-            message = (
-                'Images in subject have inconsistent shapes:'
-                f'\n{pprint.pformat(shapes_dict)}'
-            )
-            raise ValueError(message)
-
-    def check_consistent_orientation(self) -> None:
-        orientations_dict = {}
-        iterable = self.get_images_dict(intensity_only=False).items()
-        for image_name, image in iterable:
-            orientations_dict[image_name] = image.orientation
-        num_unique_orientations = len(set(orientations_dict.values()))
-        if num_unique_orientations > 1:
-            message = (
-                'Images in subject have inconsistent orientations:'
-                f'\n{pprint.pformat(orientations_dict)}'
-            )
-            raise ValueError(message)
-
     def add_transform(
             self,
             transform: 'Transform',
@@ -177,6 +167,9 @@ class Subject(dict):
         # This allows to get images using attribute notation, e.g. subject.t1
         self.__dict__.update(self)
 
-    def add_image(self, image, image_name):
+    def add_image(self, image: Image, image_name: str) -> None:
         self[image_name] = image
         self.update_attributes()
+
+    def remove_image(self, image_name: str) -> None:
+        del self[image_name]
diff --git a/torchio/transforms/augmentation/spatial/random_affine.py b/torchio/transforms/augmentation/spatial/random_affine.py
index 1ace679..d561b2d 100644
--- a/torchio/transforms/augmentation/spatial/random_affine.py
+++ b/torchio/transforms/augmentation/spatial/random_affine.py
@@ -157,7 +157,7 @@ class RandomAffine(RandomTransform, SpatialTransform):
         return transform
 
     def apply_transform(self, sample: Subject) -> dict:
-        sample.check_consistent_shape()
+        sample.check_consistent_spatial_shape()
         params = self.get_params(
             self.scales,
             self.degrees,
diff --git a/torchio/transforms/augmentation/spatial/random_elastic_deformation.py b/torchio/transforms/augmentation/spatial/random_elastic_deformation.py
index 654bf15..c6ea4a9 100644
--- a/torchio/transforms/augmentation/spatial/random_elastic_deformation.py
+++ b/torchio/transforms/augmentation/spatial/random_elastic_deformation.py
@@ -215,7 +215,7 @@ class RandomElasticDeformation(RandomTransform, SpatialTransform):
             warnings.warn(message, RuntimeWarning)
 
     def apply_transform(self, sample: Subject) -> dict:
-        sample.check_consistent_shape()
+        sample.check_consistent_spatial_shape()
         bspline_params = self.get_params(
             self.num_control_points,
             self.max_displacement,
diff --git a/torchio/transforms/preprocessing/spatial/crop_or_pad.py b/torchio/transforms/preprocessing/spatial/crop_or_pad.py
index eadaed9..a40c343 100644
--- a/torchio/transforms/preprocessing/spatial/crop_or_pad.py
+++ b/torchio/transforms/preprocessing/spatial/crop_or_pad.py
@@ -99,7 +99,7 @@ class CropOrPad(BoundsTransform):
     @staticmethod
     def _get_sample_shape(sample: Subject) -> TypeTripletInt:
         """Return the shape of the first image in the sample."""
-        sample.check_consistent_shape()
+        sample.check_consistent_spatial_shape()
         for image_dict in sample.get_images(intensity_only=False):
             data = image_dict.spatial_shape  # remove channels dimension
             break


Test patch:
diff --git a/tests/data/test_subject.py b/tests/data/test_subject.py
index 164036b..72adad4 100644
--- a/tests/data/test_subject.py
+++ b/tests/data/test_subject.py
@@ -3,6 +3,7 @@
 """Tests for Subject."""
 
 import tempfile
+import torch
 from torchio import Subject, ScalarImage, RandomFlip
 from ..utils import TorchioTestCase
 
@@ -30,3 +31,20 @@ class TestSubject(TorchioTestCase):
     def test_history(self):
         transformed = RandomFlip()(self.sample)
         self.assertIs(len(transformed.history), 1)
+
+    def test_inconsistent_shape(self):
+        subject = Subject(
+            a=ScalarImage(tensor=torch.rand(1, 2, 3, 4)),
+            b=ScalarImage(tensor=torch.rand(2, 2, 3, 4)),
+        )
+        subject.spatial_shape
+        with self.assertRaises(RuntimeError):
+            subject.shape
+
+    def test_inconsistent_spatial_shape(self):
+        subject = Subject(
+            a=ScalarImage(tensor=torch.rand(1, 3, 3, 4)),
+            b=ScalarImage(tensor=torch.rand(2, 2, 3, 4)),
+        )
+        with self.assertRaises(RuntimeError):
+            subject.spatial_shape
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 6dfb126..8a062f3 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -43,13 +43,6 @@ class TestUtils(TorchioTestCase):
         assert isinstance(guess_type('[1,3,5]'), list)
         assert isinstance(guess_type('test'), str)
 
-    def test_check_consistent_shape(self):
-        good_sample = self.sample
-        bad_sample = self.get_inconsistent_sample()
-        good_sample.check_consistent_shape()
-        with self.assertRaises(ValueError):
-            bad_sample.check_consistent_shape()
-
     def test_apply_transform_to_file(self):
         transform = RandomFlip()
         apply_transform_to_file(
diff --git a/tests/transforms/test_transforms.py b/tests/transforms/test_transforms.py
index a5c15ad..9a62306 100644
--- a/tests/transforms/test_transforms.py
+++ b/tests/transforms/test_transforms.py
@@ -51,6 +51,24 @@ class TestTransforms(TorchioTestCase):
             transforms.append(torchio.RandomLabelsToImage(label_key='label'))
         return torchio.Compose(transforms)
 
+    def test_transforms_dict(self):
+        transform = torchio.RandomNoise(keys=('t1', 't2'))
+        input_dict = {k: v.data for (k, v) in self.sample.items()}
+        transformed = transform(input_dict)
+        self.assertIsInstance(transformed, dict)
+
+    def test_transforms_dict_no_keys(self):
+        transform = torchio.RandomNoise()
+        input_dict = {k: v.data for (k, v) in self.sample.items()}
+        with self.assertRaises(RuntimeError):
+            transform(input_dict)
+
+    def test_transforms_image(self):
+        transform = self.get_transform(
+            channels=('default_image_name',), labels=False)
+        transformed = transform(self.sample.t1)
+        self.assertIsInstance(transformed, torchio.ScalarImage)
+
     def test_transforms_tensor(self):
         tensor = torch.rand(2, 4, 5, 8)
         transform = self.get_transform(
@@ -136,3 +154,9 @@ class TestTransforms(TorchioTestCase):
                 original_data,
                 f'Changes after {transform.name}'
             )
+
+
+class TestTransform(TorchioTestCase):
+    def test_abstract_transform(self):
+        with self.assertRaises(TypeError):
+            transform = torchio.Transform()