Chained affine transforms

[EdvardGrodem]

I'm trying to augment my data by doing a random rotation combined with a random translation and a random isotropic scaling. My best effort at this so far has been this:

RandAffined(keys=['image'], prob=1, rotate_range=(rot_range, rot_range, rot_range),
                    shear_range = None, translate_range=(trans_range, trans_range, trans_range),
                    padding_mode = "zeros", spatial_size=img_size, 
                    cache_grid=True, device=device,
                    ),
RandZoomd(keys=['image'], prob=1, min_zoom=scale_min, max_zoom=scale_max,  mode="area", device=device),

The issue is that this is quite slow compared to a single RandAffined. I assume that the transforms are done in stages, however, since the two transforms are affine these can theoretically be done with one stage by using their combined affine matrix. Is there a way to increase the speed by doing this in one step?

[rijobro]

Hi @EdvardGrodem, we've been keen on figuring out a way to concatenate certain transforms so as to accelerate things and also reduce the number of resampling steps (as we lose original information each time we do that). Unfortunately we haven't implemented a solution to that yet.

In your case this problem would be solved if RandAffined – which is capable of also scaling the image – could optionally randomise dimensions isotropically as well as independently. In the end it's RandAffineGrid that takes care of randomising these parameters and generating the affine transformation matrix, so we would have to enhance it here and then propagate those input arguments to all classes that use it. I can't think of a succinct way to do that other, here's my initial thought:

class RandAffineGrid(Randomizable, Transform):
    def __init__(
        self,
        rotate_range: RandRange = None,
        shear_range: RandRange = None,
        translate_range: RandRange = None,
        scale_range: RandRange = None,
        as_tensor_output: bool = True,
        device: Optional[torch.device] = None,
        rotate_independent_dims: bool = True,  # new
        shear_independent_dims: bool = True,  # new
        translate_independent_dims: bool = True,  # new
        scale_independent_dims: bool = True,  # new
    ) -> None:

Looks a little clunky to me, has anyone got any better ideas?

cc @atbenmurray (interested in concatenating multiple affine transformations).

[EdvardGrodem]

Hi @rijobro, thanks for your quick reply.
Could a solution be to change the behavior of all affine transforms if the key points to an affine matrix instead of an image? Then the transforms could accumulate by just multiplying affine matrices. When all steps are performed something line ApplyAffined could be called to actually transform the image.
Here is a use case:

LoadImaged(keys=['image']),
InitAffineMatrixd(keys=['aff_mat'], dim=3), # Initialize a identity matrix or use one from image
RandFlipd(keys=['aff_mat'], prob=0.5, spatial_axis=0), # Apply only on the matrix
RandRotate(keys=['aff_mat'], prob=1.0, range_x=0.5, range_y=0.5, range_z=0.0),
RandZoomd(keys=['aff_mat'], prob=1, min_zoom=scale_min, max_zoom=scale_max,  mode="area")
ApplyAffined(img_keys=['image'], mat_key='aff_mat) # Execute the full transformation

[rijobro]

Yes, @atbenmurray was thinking along these lines. Seems like a relatively big (but good) change to me. I think it may not be as simple as just updating the affine transformation matrix, though. For some of the transforms we have arguments such as keep_size, which would have to be passed on.

[EdvardGrodem]

Yes, I guess it is quite a big change. The shape of the image should also be stored as you said and altered with the affine matrix, so maybe these can be stored together in a VirtualImage object or something like that? This would also allow some crop transforms to be chained as well.
For now I will just write my own hacky version of RandAffine.

[rijobro]

I think in the short term, a quick hack is going to be what you need, sorry I can't offer something more eloquent.

[EdvardGrodem]

Thanks for your answer, @rijobro. I'm impressed with the community here!

[atbenmurray]

Thanks @rijobro for bringing this to my attention. @EdvardGrodem your thinking is good, and a solution that requires at most one resample is actively under discussion at present. I'll keep you informed of any progress.

In the meantime, I might suggest the following, provided you are doing your preprocessing on CPU:

1. Create a custom preprocessor class that does the building of your affine matrix using some standard homogeneous matrix manipulation library of choice. You can pretty much put all of your rng-based manipulations in there and have it return the resulting matrix.
2. Create a second custom preprocessor class that reconstructs an instance of Affine / Affined for each __call__ and then pass the matrix to it via the affine parameter before calling it internally.

Happy to discuss further if you think this approach might be helpful to you.
Best,
Ben

[EdvardGrodem]

Hi @atbenmurray!
I'm very interested in future progress on this topic!
This sounds like a good solution and some help would be really helpful. I am currently using the GPU for augmentations. Why would this solution only work for CPU?

[atbenmurray]

My guess would be that it might come with a performance implication, but it might actually be fine. There is no harm in trying it in GPU mode but keep an eye on your times.

[atbenmurray]

@EdvardGrodem Do you have enough familiarity with homogeneous transforms to get started?

[EdvardGrodem]

Yes, my background is in robotics so I'm quite familiar with homogeneous transforms. It's nice to get my hands dirty with some euler angles again. I'll publish my take on your suggestion when I've tested it :D

[EdvardGrodem]

Thanks a lot, @atbenmurray for your help. I've made a small git repo that implements your suggestion. For others that might be interested, here is a link: https://github.com/CRAI-OUS/MonaiMatrixTransforms

[atbenmurray]

This looks good. We're also moving forward with the concept of 'at most 1 resample' as an official feature request. I'll link it here if you want to track it.

[EdvardGrodem]

Yes please, that will be interesting to follow.

[wyli]

with the MetaTensor API, all the spatial transforms are tracked, and the affines updated, e.g.

MONAI/monai/transforms/spatial/array.py

Lines 1013 to 1031 in e1edc99

	self.push_transform(
	out,
	orig_size=img_t.shape[1:],
	extra_info={
	"rot_mat": transform,
	"mode": _mode,
	"padding_mode": _padding_mode,
	"align_corners": _align_corners if _align_corners is not None else TraceKeys.NONE,
	"dtype": str(_dtype)[6:], # dtype as string; remove "torch": torch.float32 -> float32
	},
	)
	return out
	def forward_meta(self, img_meta, rotate_mat):
	meta_dict = deepcopy(img_meta)
	affine = convert_data_type(img_meta["affine"], torch.Tensor, drop_meta=True)[0]
	mat = to_affine_nd(len(affine) - 1, rotate_mat)
	meta_dict["affine"] = affine @ convert_to_dst_type(mat, affine, drop_meta=True)[0]
	return meta_dict

perhaps we can leverage those APIs to implement lazy resampling -- when meta_tensor.lazy_resample is True, the transforms will only update meta_tensor.applied_operations attributes to track the transforms. and when the user set meta_tensor.lazy_resample=False, actual voxels values will be updated...

class MetaTensor:
    
    @property
    def lazy_resample() -> bool:
        return self._lazy_resample
    
    @lazy_resample.setter
    def lazy_resample(self, flag: bool):
        if not flag and self._lazy_resample:  # previously lazy_resample=True, and now
            self.do_spatial_resampling(self.applied_operations)
        return

[atbenmurray]

Hi @wyli, yes, @rijobro and I had a conversation about metatensor and lazy resampling. I'm in the middle of a WIP implementation; Richard suggested that I implement it first and then we look at how we might consolidate the concepts. That said, it looks like the API is heading towards maturity and it might be better for me to implement it with that more in mind.

[wyli]

yes, metatensor functionality wise is in a good shape, but the API still needs some clean up. possibly there are requirements out of this thread we should take into consideration for the metatensor API design...