Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add combined affine and elastic deformation augmentation #1073

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

mrdkucher
Copy link

@mrdkucher mrdkucher commented Apr 19, 2023

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

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

@mrdkucher mrdkucher marked this pull request as ready for review April 19, 2023 20:01
@mrdkucher mrdkucher force-pushed the 1052-add-combined-affine-elastic-augmentation branch from 23dbc31 to b3559b7 Compare April 20, 2023 15:46
@mrdkucher
Copy link
Author

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

Copy link
Contributor

@romainVala romainVala left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@mrdkucher
Copy link
Author

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
Copy link
Contributor

romainVala commented May 16, 2023

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
Copy link
Owner

fepegar commented May 16, 2023

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

@romainVala
Copy link
Contributor

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
Copy link
Author

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
Copy link
Owner

fepegar commented Jun 5, 2023

@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
Copy link
Contributor

@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
Copy link
Contributor

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
Copy link
Owner

fepegar commented Jun 6, 2023

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

@mrdkucher
Copy link
Author

@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 mrdkucher force-pushed the 1052-add-combined-affine-elastic-augmentation branch from b3559b7 to 2bf9a89 Compare June 9, 2023 03:03
@mrdkucher
Copy link
Author

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:
Screenshot 2023-06-08 at 9 37 54 PM

@mrdkucher
Copy link
Author

@fepegar can we merge this?

@mrdkucher
Copy link
Author

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

@fepegar
Copy link
Owner

fepegar commented Aug 3, 2023

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
Copy link
Author

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

SimpleITK's ResampleImageFilter is an expensive operation,
especially when its 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.

Add combined affine and elastic deformation augmentation

SimpleITK's ResampleImageFilter is an expensive operation,
especially when its 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.

Resolves: fepegar#1052

Remove duplicated code

Update pytests
@mrdkucher mrdkucher force-pushed the 1052-add-combined-affine-elastic-augmentation branch from 2ad1090 to 0d7bee3 Compare August 27, 2023 22:26
@mrdkucher
Copy link
Author

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

@mrdkucher
Copy link
Author

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

@fepegar
Copy link
Owner

fepegar commented Sep 28, 2023

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Combined Affine/Elastic augmentations
3 participants