From 2ad1090a92feffb2a4c26bd559aea536308903d2 Mon Sep 17 00:00:00 2001 From: David Kucher Date: Fri, 25 Aug 2023 14:37:15 -0700 Subject: [PATCH] Update pytests --- .../random_affine_elastic_deformation.py | 57 ++-- .../test_random_affine_elastic_deformation.py | 253 +++++++++++++----- tests/transforms/test_transforms.py | 2 +- 3 files changed, 215 insertions(+), 97 deletions(-) diff --git a/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py b/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py index aa520515..1f62965a 100644 --- a/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py +++ b/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py @@ -1,22 +1,21 @@ -from typing import Any, Dict, Optional, Union, Tuple +from typing import Any +from typing import Dict +from typing import Optional import numpy as np import SimpleITK as sitk import torch -from .random_affine import Affine, RandomAffine -from .random_elastic_deformation import ElasticDeformation, RandomElasticDeformation +from .random_affine import Affine +from .random_affine import RandomAffine +from .random_elastic_deformation import ElasticDeformation +from .random_elastic_deformation import RandomElasticDeformation from .. import RandomTransform from ... import SpatialTransform from ....constants import INTENSITY from ....constants import TYPE from ....data.io import nib_to_sitk from ....data.subject import Subject -from ....typing import TypeRangeFloat -from ....typing import TypeSextetFloat -from ....typing import TypeTripletFloat - -TypeOneToSixFloat = Union[TypeRangeFloat, TypeTripletFloat, TypeSextetFloat] class RandomCombinedAffineElasticDeformation(RandomTransform, SpatialTransform): @@ -90,12 +89,29 @@ def get_params(self): def apply_transform(self, subject: Subject): affine_params, elastic_params = self.get_params() + scaling_params, rotation_params, translation_params = affine_params + affine_params = { + 'scales': scaling_params.tolist(), + 'degrees': rotation_params.tolist(), + 'translation': translation_params.tolist(), + 'center': self.random_affine.center, + 'default_pad_value': self.random_affine.default_pad_value, + 'image_interpolation': self.random_affine.image_interpolation, + 'label_interpolation': self.random_affine.label_interpolation, + 'check_shape': self.random_affine.check_shape, + } + + elastic_params = { + 'control_points': elastic_params, + 'max_displacement': self.random_elastic.max_displacement, + 'image_interpolation': self.random_elastic.image_interpolation, + 'label_interpolation': self.random_elastic.label_interpolation, + } + arguments = { 'affine_first': self.affine_first, 'affine_params': affine_params, 'elastic_params': elastic_params, - 'affine_kwargs': self.affine_kwargs, - 'elastic_kwargs': self.elastic_kwargs, } transform = CombinedAffineElasticDeformation( @@ -125,32 +141,25 @@ class CombinedAffineElasticDeformation(SpatialTransform): def __init__( self, affine_first: bool, - affine_params: Tuple[TypeTripletFloat], - elastic_params: np.ndarray, - affine_kwargs: Dict[str, Any], - elastic_kwargs: Dict[str, Any], + affine_params: Dict[str, Any], + elastic_params: Dict[str, Any], **kwargs, ) -> None: super().__init__(**kwargs) self.affine_first = affine_first + self.affine_params = affine_params self._affine = Affine( - *affine_params**affine_kwargs, + **self.affine_params, **kwargs, ) + self.elastic_params = elastic_params self._elastic = ElasticDeformation( - *elastic_params, - **elastic_kwargs, + **self.elastic_params, **kwargs, ) - self.args_names = [ - 'affine_params', - 'elastic_params', - 'affine_first', - 'affine_kwargs', - 'elastic_kwargs', - ] + self.args_names = ['affine_first', 'affine_params', 'elastic_params'] def apply_transform(self, subject: Subject) -> Subject: if self._affine.check_shape: diff --git a/tests/transforms/augmentation/test_random_affine_elastic_deformation.py b/tests/transforms/augmentation/test_random_affine_elastic_deformation.py index df11f635..cc138d73 100644 --- a/tests/transforms/augmentation/test_random_affine_elastic_deformation.py +++ b/tests/transforms/augmentation/test_random_affine_elastic_deformation.py @@ -32,29 +32,38 @@ def test_inputs_interpolation(self): def test_num_control_points_noint(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(num_control_points=2.5) + tio.RandomCombinedAffineElasticDeformation( + elastic_kwargs={'num_control_points': 2.5} + ) def test_num_control_points_small(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(num_control_points=3) + tio.RandomCombinedAffineElasticDeformation( + elastic_kwargs={'num_control_points': 3} + ) def test_max_displacement_no_num(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(max_displacement=None) + tio.RandomCombinedAffineElasticDeformation( + elastic_kwargs={'max_displacement': None} + ) def test_max_displacement_negative(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(max_displacement=-1) + tio.RandomCombinedAffineElasticDeformation( + elastic_kwargs={'max_displacement': -1} + ) def test_wrong_locked_borders(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(locked_borders=-1) + tio.RandomCombinedAffineElasticDeformation( + elastic_kwargs={'locked_borders': -1} + ) def test_coarse_grid_removed(self): with pytest.raises(ValueError): tio.RandomCombinedAffineElasticDeformation( - num_control_points=(4, 5, 6), - locked_borders=2, + elastic_kwargs={'num_control_points': (4, 5, 6), 'locked_borders': 2} ) def test_folding(self): @@ -62,23 +71,31 @@ def test_folding(self): # Then grid spacing is (10/(12-2), 20/(5-2), 30/(5-2)) # or (1, 6.7, 10), and half is (0.5, 3.3, 5) transform = tio.RandomCombinedAffineElasticDeformation( - num_control_points=(12, 5, 5), - max_displacement=6, + elastic_kwargs={'num_control_points': (12, 5, 5), 'max_displacement': 6} ) with pytest.warns(RuntimeWarning): transform(self.sample_subject) def test_num_control_points(self): - tio.RandomCombinedAffineElasticDeformation(num_control_points=5) - tio.RandomCombinedAffineElasticDeformation(num_control_points=(5, 6, 7)) + tio.RandomCombinedAffineElasticDeformation( + elastic_kwargs={'num_control_points': 5} + ) + tio.RandomCombinedAffineElasticDeformation( + elastic_kwargs={'num_control_points': (5, 6, 7)} + ) def test_max_displacement(self): - tio.RandomCombinedAffineElasticDeformation(max_displacement=5) - tio.RandomCombinedAffineElasticDeformation(max_displacement=(5, 6, 7)) + tio.RandomCombinedAffineElasticDeformation( + elastic_kwargs={'max_displacement': 5} + ) + tio.RandomCombinedAffineElasticDeformation( + elastic_kwargs={'max_displacement': (5, 6, 7)} + ) def test_no_displacement(self): transform = tio.RandomCombinedAffineElasticDeformation( - max_displacement=0, scales=0, degrees=0, translation=0 + affine_kwargs={'scales': 0, 'degrees': 0, 'translation': 0}, + elastic_kwargs={'max_displacement': 0}, ) transformed = transform(self.sample_subject) self.assert_tensor_equal( @@ -93,9 +110,11 @@ def test_no_displacement(self): def test_rotation_image(self): # Rotation around image center transform = tio.RandomCombinedAffineElasticDeformation( - degrees=(90, 90), - default_pad_value=0, - center='image', + affine_kwargs={ + 'degrees': (90, 90), + 'default_pad_value': 0, + 'center': 'image', + } ) transformed = transform(self.sample_subject) total = transformed.t1.data.sum() @@ -104,9 +123,11 @@ def test_rotation_image(self): def test_rotation_origin(self): # Rotation around far away point, image should be empty transform = tio.RandomCombinedAffineElasticDeformation( - degrees=(90, 90), - default_pad_value=0, - center='origin', + affine_kwargs={ + 'degrees': (90, 90), + 'default_pad_value': 0, + 'center': 'origin', + } ) transformed = transform(self.sample_subject) total = transformed.t1.data.sum() @@ -114,24 +135,27 @@ def test_rotation_origin(self): def test_no_rotation(self): transform = tio.RandomCombinedAffineElasticDeformation( - scales=(1, 1), - degrees=(0, 0), - default_pad_value=0, - max_displacement=0, - center='image', + affine_kwargs={ + 'scales': (1, 1), + 'degrees': (0, 0), + 'default_pad_value': 0, + 'center': 'image', + }, + elastic_kwargs={'max_displacement': 0}, ) transformed = transform(self.sample_subject) self.assert_tensor_almost_equal( self.sample_subject.t1.data, transformed.t1.data, ) - transform = tio.RandomCombinedAffineElasticDeformation( - scales=(1, 1), - degrees=(180, 180), - default_pad_value=0, - max_displacement=0, - center='image', + affine_kwargs={ + 'scales': (1, 1), + 'degrees': (180, 180), + 'default_pad_value': 0, + 'center': 'image', + }, + elastic_kwargs={'max_displacement': 0}, ) transformed = transform(self.sample_subject) transformed = transform(transformed) @@ -141,105 +165,169 @@ def test_no_rotation(self): ) def test_isotropic(self): - tio.RandomCombinedAffineElasticDeformation(isotropic=True)(self.sample_subject) - - def test_mean(self): - tio.RandomCombinedAffineElasticDeformation(default_pad_value='mean')( + tio.RandomCombinedAffineElasticDeformation(affine_kwargs={'isotropic': True})( self.sample_subject ) + def test_mean(self): + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'default_pad_value': 'mean'} + )(self.sample_subject) + def test_otsu(self): - tio.RandomCombinedAffineElasticDeformation(default_pad_value='otsu')( - self.sample_subject - ) + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'default_pad_value': 'otsu'} + )(self.sample_subject) def test_bad_center(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(center='bad') + tio.RandomCombinedAffineElasticDeformation(affine_kwargs={'center': 'bad'}) def test_negative_scales(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(scales=(-1, 1)) + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'scales': (-1, 1)} + ) def test_scale_too_large(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(scales=1.5) + tio.RandomCombinedAffineElasticDeformation(affine_kwargs={'scales': 1.5}) def test_scales_range_with_negative_min(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(scales=(-1, 4)) + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'scales': (-1, 4)} + ) def test_wrong_scales_type(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(scales='wrong') + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'scales': 'wrong'} + ) def test_wrong_degrees_type(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(degrees='wrong') + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'degrees': 'wrong'} + ) def test_too_many_translation_values(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(translation=(-10, 4, 42)) + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'translation': (-10, 4, 42)} + ) def test_wrong_translation_type(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(translation='wrong') + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'translation': 'wrong'} + ) def test_wrong_center(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(center=0) + tio.RandomCombinedAffineElasticDeformation(affine_kwargs={'center': 0}) def test_wrong_default_pad_value(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(default_pad_value='wrong') + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'default_pad_value': 'wrong'} + ) def test_wrong_image_interpolation_type(self): with pytest.raises(TypeError): - tio.RandomCombinedAffineElasticDeformation(image_interpolation=0) + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'image_interpolation': 0} + ) def test_wrong_image_interpolation_value(self): with pytest.raises(ValueError): - tio.RandomCombinedAffineElasticDeformation(image_interpolation='wrong') + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'image_interpolation': 'wrong'} + ) def test_incompatible_args_isotropic(self): with pytest.raises(ValueError): tio.RandomCombinedAffineElasticDeformation( - scales=(0.8, 0.5, 0.1), isotropic=True + affine_kwargs={'scales': (0.8, 0.5, 0.1), 'isotropic': True} ) def test_parse_scales(self): def do_assert(transform): - assert transform.scales == 3 * (0.9, 1.1) + assert transform.random_affine.scales == 3 * (0.9, 1.1) - do_assert(tio.RandomCombinedAffineElasticDeformation(scales=0.1)) - do_assert(tio.RandomCombinedAffineElasticDeformation(scales=(0.9, 1.1))) - do_assert(tio.RandomCombinedAffineElasticDeformation(scales=3 * (0.1,))) - do_assert(tio.RandomCombinedAffineElasticDeformation(scales=3 * [0.9, 1.1])) + do_assert( + tio.RandomCombinedAffineElasticDeformation(affine_kwargs={'scales': 0.1}) + ) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'scales': (0.9, 1.1)} + ) + ) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'scales': 3 * (0.1,)} + ) + ) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'scales': 3 * [0.9, 1.1]} + ) + ) def test_parse_degrees(self): def do_assert(transform): - assert transform.degrees == 3 * (-10, 10) + assert transform.random_affine.degrees == 3 * (-10, 10) - do_assert(tio.RandomCombinedAffineElasticDeformation(degrees=10)) - do_assert(tio.RandomCombinedAffineElasticDeformation(degrees=(-10, 10))) - do_assert(tio.RandomCombinedAffineElasticDeformation(degrees=3 * (10,))) - do_assert(tio.RandomCombinedAffineElasticDeformation(degrees=3 * [-10, 10])) + do_assert( + tio.RandomCombinedAffineElasticDeformation(affine_kwargs={'degrees': 10}) + ) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'degrees': (-10, 10)} + ) + ) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'degrees': 3 * (10,)} + ) + ) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'degrees': 3 * [-10, 10]} + ) + ) def test_parse_translation(self): def do_assert(transform): - assert transform.translation == 3 * (-10, 10) + assert transform.random_affine.translation == 3 * (-10, 10) - do_assert(tio.RandomCombinedAffineElasticDeformation(translation=10)) - do_assert(tio.RandomCombinedAffineElasticDeformation(translation=(-10, 10))) - do_assert(tio.RandomCombinedAffineElasticDeformation(translation=3 * (10,))) - do_assert(tio.RandomCombinedAffineElasticDeformation(translation=3 * [-10, 10])) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'translation': 10} + ) + ) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'translation': (-10, 10)} + ) + ) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'translation': 3 * (10,)} + ) + ) + do_assert( + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'translation': 3 * [-10, 10]} + ) + ) def test_default_value_label_map(self): # From https://github.com/fepegar/torchio/issues/626 a = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).reshape(1, 3, 3, 1) image = tio.LabelMap(tensor=a) aff = tio.RandomCombinedAffineElasticDeformation( - translation=(0, 1, 1), default_pad_value='otsu' + affine_kwargs={'translation': (0, 1, 1), 'default_pad_value': 'otsu'} ) transformed = aff(image) assert all(n in (0, 1) for n in transformed.data.flatten()) @@ -266,22 +354,43 @@ def test_different_spaces(self): new_subject = tio.Subject(t1=t1, label=label) with pytest.raises(RuntimeError): tio.RandomCombinedAffineElasticDeformation()(new_subject) - tio.RandomCombinedAffineElasticDeformation(check_shape=False)(new_subject) + tio.RandomCombinedAffineElasticDeformation( + affine_kwargs={'check_shape': False} + )(new_subject) def test_transform_order(self): src_transform = tio.RandomCombinedAffineElasticDeformation( - scales=0, degrees=0, translation=1, num_control_points=5, max_displacement=1 + affine_kwargs={'scales': 0, 'degrees': 0, 'translation': 1}, + elastic_kwargs={'num_control_points': 5, 'max_displacement': 1}, ) (scales, degrees, translation), control_points = src_transform.get_params() - max_displacement = src_transform.max_displacement + max_displacement = src_transform.random_elastic.max_displacement transform1 = tio.CombinedAffineElasticDeformation( - True, scales, degrees, translation, control_points, max_displacement + affine_first=True, + affine_params={ + 'scales': scales, + 'degrees': degrees, + 'translation': translation, + }, + elastic_params={ + 'control_points': control_points, + 'max_displacement': max_displacement, + }, ) transform2 = tio.CombinedAffineElasticDeformation( - False, scales, degrees, translation, control_points, max_displacement + affine_first=False, + affine_params={ + 'scales': scales, + 'degrees': degrees, + 'translation': translation, + }, + elastic_params={ + 'control_points': control_points, + 'max_displacement': max_displacement, + }, ) transformed1 = transform1(self.sample_subject) diff --git a/tests/transforms/test_transforms.py b/tests/transforms/test_transforms.py index 59111be5..55c0a0de 100644 --- a/tests/transforms/test_transforms.py +++ b/tests/transforms/test_transforms.py @@ -18,7 +18,7 @@ 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) affine_elastic = tio.RandomCombinedAffineElasticDeformation( - max_displacement=disp + elastic_kwargs={'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)