RandAffined fails with KeyError when allow_missing_keys and key at index 0 is missing

[Spenhouet]

Describe the bug
We have a pipeline in which not all samples have the same keys.
The RandAffined transform accesses fallback data via index 0.
If index 0 is a missing key, then this transform will fail with a KeyError.

805   sp_size = fall_back_tuple(self.rand_affine.spatial_size, data[self.keys[0]].shape[1:])
806   # change image size or do random transform
807   do_resampling = self._do_transform or (sp_size != ensure_tuple(data[self.keys[0]].shape[1:]))

Expected behavior
If the option allow_missing_keys is set, I would expect the RandAffined transform to take the first element of the existing keys as fallback.

[Nic-Ma]

Hi @Spenhouet ,

Thanks for your feedback.
It sounds like a good enhancement for the components, we have the logic data[self.keys[0]].shape[1:] in several transforms. I will try to update them based on your idea.

Thanks.

[Spenhouet]

@Nic-Ma Any idea how I could work around this without needing to patch all augmentation transforms?
I tried adding a fake image as first key but this is then actually used to resize the image (in the RandAffine transform). Probably this is an independent bug report.

[Spenhouet]

Okay, the RandAffine has much more problems. It is not working at all for us... I will use the torch.io transform instead. I don't have the time to open the bug report currently. I have probably more than 20 bug reports on my list which I did not yet report. It feels like I have half of MONAI patched in my code base... this really start to get frustrating.

[Nic-Ma]

Hi @Spenhouet ,

Thanks for your great feedback.
Please feel free to develop your project in your preferred method first, we will try to enhance MONAI APIs ASAP.

Thanks.

[Spenhouet]

Just to note down: RandGaussianNoised has the same issue with missing keys.

[Nic-Ma]

Hi @Spenhouet ,

Thanks for your feedback.
I am already working on a PR to address this issue, will submit for review soon.

Thanks.

[Nic-Ma]

Hi @Spenhouet ,

The PR had been merged into dev to fix this issue.

Thanks for your feedback.