Efficient spatial augmentation without edge effects after RandCropByPosNegLabeld

[maunzzz]

Hi!
I'm working on a 3d segmentation problem with large CT images (512x512x400).
I want to do spatial augmentation in my training pipeline. For performance issues I want this to be done after the RandCropByPosNegLabeld transform so I don't need to apply the random affine transform on the entire image.
However, if I do RandCropByPosNegLabeld with my input size (96,96,96) and then apply a random affine transformation (as in the transform example below, Alternative 1) there will be unnecessary edge effects.

To avoid this, I could extract a larger (large enough to guarantee that no "outside" pixels will be transformed into the center 96x96x96) patch with RandCropByPosNegLabeld and then apply a random affine transform to the larger patch, as in Alternative 2 below.

Non of these transform pipelines are optimal, in alternative 1 we get unnecessary edge effects and in alternative 2 we apply an affine transform to an unnecessary large patch. Is there a better way of doing this in Monai?

Alternative 1

transforms = Compose([LoadImaged(keys=['CT', 'Segmentation']),
                      EnsureChannelFirstd(keys=['CT', 'Segmentation']),
                      Spacingd(keys=['CT', 'Segmentation'], pixdim=(1.0, 1.0, 1.0), mode=('bilinear', 'nearest')),
                      SpatialPadd(keys=['CT', 'Segmentation'], spatial_size=(96, 96, 96), method=Method.SYMMETRIC,
                                  mode=NumpyPadMode.CONSTANT, constant_values=-1024),
                      ScaleIntensityRanged(keys='CT', a_min=-1024, a_max=3000, b_min=-1, b_max=1, clip=True),
                      RandCropByPosNegLabeld(keys=['CT', 'Segmentation'],
                                             label_key='Segmentation',
                                             spatial_size=(96, 96, 96),
                                             pos=1,
                                             neg=1,
                                             num_samples=4,
                                             allow_smaller=False),

                      RandAffined(keys=['CT', 'Segmentation'],
                                  spatial_size=(96, 96, 96),
                                  rotate_range=(-np.pi, np.pi),
                                  scale_range=(-0.1, 0.1))])

Alternative 2

transforms = Compose([LoadImaged(keys=['CT', 'Segmentation']),
                      EnsureChannelFirstd(keys=['CT', 'Segmentation']),
                      Spacingd(keys=['CT', 'Segmentation'], pixdim=(1.0, 1.0, 1.0), mode=('bilinear', 'nearest')),
                      SpatialPadd(keys=['CT', 'Segmentation'], spatial_size=(96, 96, 96), method=Method.SYMMETRIC,
                                  mode=NumpyPadMode.CONSTANT, constant_values=-1024),
                      ScaleIntensityRanged(keys='CT', a_min=-1024, a_max=3000, b_min=-1, b_max=1, clip=True),
                      RandCropByPosNegLabeld(keys=['CT', 'Segmentation'],
                                             label_key='Segmentation',
                                             spatial_size=(150, 150, 150),
                                             pos=1,
                                             neg=1,
                                             num_samples=4,
                                             allow_smaller=False),

                      RandAffined(keys=['CT', 'Segmentation'],
                                  spatial_size=(96, 96, 96),
                                  rotate_range=(-np.pi, np.pi),
                                  scale_range=(-0.1, 0.1))])

Thanks!

[KumoLiu]

Hi @maunzzz, thanks for your interest here. Maybe this feature can help solve the issue here which will come out soon. The feature is still in develop and will have an initial version in the later release. With lazy resampling, we could increase the transform speed, reduce memory and artifact reduction. When it is finished, welcome to test out this feature in your pipeline as well. Thanks!

[maunzzz]

Hi @KumoLiu, thanks for your answer. The lazy evaluation feature looks interesting but does it really help with my problem here?
The problem is not really that I apply a set of affine transforms and resamplings that could be replaced by a single affine transform. The problem is that I need to do a cropping in RandCropByPosNegLabeld before the random affine transformation.

[KumoLiu]

Hi @maunzzz, yes I know that. The problem you said is that if you crop first, then there will be unnecessary edge effects. But if using lazy resampling, you may don't need to crop first. Thanks!

[maunzzz]

Alright, then I understand. Do you know if there are any better way of doing this with the code currently available in monai?

[wyli]

@maunzzz these two options are currently what we have, there's a cache_grid option for RandAffine to save some cycles (but with more space)...

[maunzzz]

Thanks for the answers and interest in this thread.
I solved it by implementing my own class that combines RandCropByPosNegLabeld and RandAffined. Instead of having RandCropByPosNegLabeld cropping a part of the 3d volume the class applies a random affine transform with the wanted patch spatial_size.