Add combined affine and elastic deformation augmentation

[mrdkucher]

Fixes #1052.

Description

SimpleITK's ResampleImageFilter is an expensive operation, especially when it's called sequentially for both affine and elastic transformations. By combining the SimpleITK transforms for both augmentations, the processing time for a sample can be reduced.

Currently the augmentation uses a combined probability for applying the transform, rather than independently applying them.

Timing 128^3 Combined:

python -m timeit -n 3 -r 3 -s "import torch; from torchio import LabelMap, ScalarImage, Subject, RandomCombinedAffineElasticDeformation; image_size = 128; volume = ScalarImage(tensor=torch.randn(1, image_size, image_size, image_size)); label = LabelMap(tensor=torch.randint(10, (1, image_size, image_size, image_size))); subject = Subject(volume=volume, label=label); transform = RandomCombinedAffineElasticDeformation(scales=[0.15, 0.15, 0.15], degrees=[0., 0., 25.], default_pad_value=0., num_control_points=[7, 7, 7], max_displacement=[6, 6, 27])" "transform(subject)"
3 loops, best of 3: 402 msec per loop

128^3 Sequential:

python -m timeit -n 3 -r 3 -s "import torch; from torchio import LabelMap, ScalarImage, Subject, RandomAffine, RandomElasticDeformation, Compose; image_size = 128; volume = ScalarImage(tensor=torch.randn(1, image_size, image_size, image_size)); label = LabelMap(tensor=torch.randint(10, (1, image_size, image_size, image_size))); subject = Subject(volume=volume, label=label); transform = Compose([RandomAffine(scales=[0.15, 0.15, 0.15], degrees=[0., 0., 25.], default_pad_value=0.), RandomElasticDeformation(num_control_points=[7, 7, 7], max_displacement=[6, 6, 27])])" "transform(subject)" 
3 loops, best of 3: 480 msec per loop

512^3 Combined: 3 loops, best of 3: 22.1 sec per loop (avg 7.4 s per call) 512^3 Sequential: 3 loops, best of 3: 34.8 sec per loop (avg 11.6 per call)

Checklist

• [X] I have read the CONTRIBUTING docs and have a developer setup (especially important are pre-commitand pytest)
• [X] Non-breaking change (would not break existing functionality)
• [ ] Breaking change (would cause existing functionality to change)
• [X] Tests added or modified to cover the changes
• [X] Integration tests passed locally by running pytest
• [X] In-line docstrings updated
• [X] Documentation updated, tested running make html inside the docs/ folder
• [X] This pull request is ready to be reviewed

[mrdkucher]

@fepegar Could you review this when you have a chance? Thanks!

[mrdkucher]

Sorry for the late delay but I do think this is a nice improvement

I may not be the best to review code, but It looks nice to me

Thanks for the review @romainVala! Could you start the workflows for testing this branch?

[romainVala]

Hello sorry for my "non professional" review, I first just look at the code But you are right: it is more important to test it,

I just tried it and get an issue with sitk version Unfortunatelly I work with sitk version 1

from torchio.utils import get_major_sitk_version
get_major_sitk_version()
Out[11]: 1

This is why I get an error in random_affine_elastic_deformation.py line 407 (module 'SimpleITK' has no attribute 'CompositeTransform')

This is handel in the RandomAffine transform (line 328), but may be it will be harder to integrate here ?

[fepegar]

Thanks for the PR, @mrdkucher, and @romainVala for reviewing. I'll try to take a look this week.

[romainVala]

Thanks for the PR, @mrdkucher, and @romainVala for reviewing. I'll try to take a look this week.

what do you think about this sitk version ? should torchio stay with a backward compatibility ?

The reason I use stik major version 1 is because I install SimpleElastix (https://simpleelastix.github.io/) instead of SimpleItk (in order to have coregistration functionalities) ... and I do not see the version 2 this install is quite weird because it is the same name as the SimpleItk (I do not understand why it is so)

I understand this is quite borderline (this is quite a mess,) and I need to see how to change it, but just curious to know if I need to do it quick or if this PR can easily accomodate with previous stik version

[mrdkucher]

Thanks for the PR, @mrdkucher, and @romainVala for reviewing. I'll try to take a look this week.

what do you think about this sitk version ? should torchio stay with a backward compatibility ?

The reason I use stik major version 1 is because I install SimpleElastix (https://simpleelastix.github.io/) instead of SimpleItk (in order to have coregistration functionalities) ... and I do not see the version 2 this install is quite weird because it is the same name as the SimpleItk (I do not understand why it is so)

I understand this is quite borderline (this is quite a mess,) and I need to see how to change it, but just curious to know if I need to do it quick or if this PR can easily accomodate with previous stik version

Thanks for the review @fepegar! I think there's a different API for composed transforms in sitk v1, so I'll try to incorporate it as well.

[fepegar]

@mrdkucher Thanks for the PR! I think it generally looks good. I might make a couple of minor style changes. My main concern is the copied docs. Shall we instead add links to RandomAffine and RandomElasticDeformation for an explanation of the parameters? Also, can you please take a look at that mypy error?

@romainVala, SimpleITK 2 was released almost three years ago. I think we can merge this and, if you need to use the transform, you could work on adding compatibility with older versions of SimpleITK. Does that make sense?

[romainVala]

@romainVala, SimpleITK 2 was released almost three years ago. I think we can merge this and, if you need to use the transform, you could work on adding compatibility with older versions of SimpleITK. Does that make sense? Sure no problem for me, I'll adapt

I do use old version because I need SimpleElastix (which I do not know why is installed with the same name as simpleitk (but with the old version)

[romainVala]

do I need this transform is an interesting question ... except for the gain in computation time, does it make any difference to have only 1 reslicing operation instead of 2 ?

[fepegar]

I suppose there's also an advantage in interpolating only once.

[mrdkucher]

@mrdkucher Thanks for the PR! I think it generally looks good. I might make a couple of minor style changes. My main concern is the copied docs. Shall we instead add links to RandomAffine and RandomElasticDeformation for an explanation of the parameters? Also, can you please take a look at that mypy error?

@fepegar Absolutely, let me make the changes with just links, and push a final version with the mypy issue resolved.

[mrdkucher]

review this please @fepegar. I resolved the mypy issues using the same workaround present here: https://github.com/fepegar/torchio/blob/main/src/torchio/transforms/augmentation/spatial/random_elastic_deformation.py#L305.

I also verified the docs are built correctly and the links work for the argument explanations. Here's a screenshot:

[mrdkucher]

@fepegar can we merge this?

[mrdkucher]

Checking in again @fepegar @romainVala, can we merge this PR?

[fepegar]

I've added some minor style changes and changed the example for the docs.

My concern with this implementation is all the copy & pasted code from the two transforms, which makes things difficult to maintain. Is there a way to do something like

class YourNewTransform(...):
    def __init__(self, affine_kwargs, elastic_kwargs):
        self.affine_transform = RandomAffine(**affine_kwargs)
        self.elastic_transform = RandomElastic(**elastic_kwargs)

That way, all the defaults, checks and parsing are handled by the specialised transforms, and we avoid duplication.

[mrdkucher]

@fepegar Thanks for the suggestion. Took a bit of wrangling, but this removes lots of duplicated code.

[mrdkucher]

@fepegar Could I get another review when you have a chance?

[mrdkucher]

@fepegar If you have a moment, could I get a review?

[fepegar]

Hi, @mrdkucher. I'll try to take a look as soon as I can. You can use your branch in the meantime, if you need the transform urgently.