Improve coverage

fepegar · Sep 1, 2021 · 2a9c091 · 2a9c091
1 parent 761aac1
commit 2a9c091
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/tests/transforms/preprocessing/test_resize.py b/tests/transforms/preprocessing/test_resize.py
@@ -1,3 +1,5 @@
+import torch
+import numpy as np
 import torchio as tio
 from ...utils import TorchioTestCase
 
@@ -17,3 +19,13 @@ def test_all_dims(self):
         transformed = transform(self.sample_subject)
         for image in transformed.get_images(intensity_only=False):
             self.assertEqual(image.spatial_shape, target_shape)
+
+    def test_fix_shape(self):
+        # We use values that are known to need cropping
+        tensor = torch.rand(1, 8, 180, 320)
+        affine = np.diag((5, 1, 1, 1))
+        im = tio.ScalarImage(tensor=tensor, affine=affine)
+        target = 12
+        with self.assertWarns(UserWarning):
+            result = tio.Resize(target)(im)
+        self.assertEqual(result.spatial_shape, 3 * (target,))
diff --git a/tests/transforms/test_transforms.py b/tests/transforms/test_transforms.py
@@ -17,16 +17,18 @@ def get_transform(self, channels, is_3d=True, labels=True):
         disp = 1 if is_3d else (1, 1, 0.01)
         elastic = tio.RandomElasticDeformation(max_displacement=disp)
         cp_args = (9, 21, 30) if is_3d else (21, 30, 1)
+        resize_args = (10, 20, 30) if is_3d else (10, 20, 1)
         flip_axes = axes_downsample = (0, 1, 2) if is_3d else (0, 1)
         swap_patch = (2, 3, 4) if is_3d else (3, 4, 1)
         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)
         remapping = {1: 2, 2: 1, 3: 20, 4: 25}
         transforms = [
             tio.CropOrPad(cp_args),
+            tio.EnsureShapeMultiple(2, method='crop'),
+            tio.Resize(resize_args),
             tio.ToCanonical(),
             tio.RandomAnisotropy(downsampling=(1.75, 2), axes=axes_downsample),
-            tio.EnsureShapeMultiple(2, method='crop'),
             tio.CopyAffine(channels[0]),
             tio.Resample((1, 1.1, 1.25)),
             tio.RandomFlip(axes=flip_axes, flip_probability=1),