Repository: Project-MONAI/MONAI

Issue:
Random Intensity transforms for test-time augmentation
Hi,

Test-time augmentation is a helpful tool to use but it seems that currently only spatial transformations are included for TTA. Intensity transformations such as Random contrast are not invertible transforms and thus cannot be used in current TTA implementation. Is there a way that I can use intensity transforms as TTA? Thanks.
Random Intensity transforms for test-time augmentation
Hi,

Test-time augmentation is a helpful tool to use but it seems that currently only spatial transformations are included for TTA. Intensity transformations such as Random contrast are not invertible transforms and thus cannot be used in current TTA implementation. Is there a way that I can use intensity transforms as TTA? Thanks.


Discussion:
Hi @rijobro ,

Could you please help share some comments about the TTA plan here?

Thanks in advance.

Hi @han-liu we certainly have it on our agenda to extend invertible transforms to include more than just spatial ones. 
Which intensity transform in particular are you using? Perhaps I can make this one invertible to be used as an example for future intensity transforms.
I guess invertible transforms don't have to be necessarilly mathematically invertible. We just need to define a suitable pseudo-inverse for it.

The example of adding Gaussian noise on the image intensities is interesting to consider. We could either make it an invertible transform by storing the noise added at test-time or just forget about the noise values and consider that the pseudo-inverse is the identity mapping. In any case, if TTA with Gaussian noise on the image intensities is to be used for segmentation purposes, the noise values are irrelevant (just a nuisance paraneters) for the segmentation output. In this case the pseudo-inverse to be applied to the network output is the identity.

Do we already handle different specialisation of the inverse transformation based on the type of network output?

In terms of really non-invertible transforms (those where storing the parameters is not enough to revert it up to numerical issues), I don't have a realistic example in mind. However, I guess that considering an image intensity change transform whereby we take the absolute value of the input intensity is a sufficient test case. In this scenario, we might want to have no defined inverse if the network is an auto-encoder but consider an identity pseudo-inverse if the network output is a segmentation map.

By the way, this raise the question of what to do for networks which are note image-in-image-out, e.g. classifiers. There as well, we might need to specialise the definition of inverse.

So in short, we should probably include some knowledge of the codomain and whether the transform acts of the input domain and output codomain in an [equivariant](https://en.wikipedia.org/wiki/Equivariant_map) fashion.
@rijobro @tvercaut Thanks for the clarification and I think what @tvercaut proposed here is a great feature to be added to TTA.
@tvercaut to take the example of `GaussianNoise`, the inverse of the spatial component of that transform is identity, because obviously adding Gaussian noise doesn't modify the image spatially. Having said that, I think it would be complicated for us to therefore say that a pseudo of the inverse can be given as identity. Because at that point, we're saying that the most important component of the inverse is the spatial component. This may or may not be true, but I'm not sure we really want to make a decision either way. 

A potential solution might be to have something like:

```python
def inverse(self, x):
    x = self.spatial_inverse(x)
    x = self.intensity_inverse(x)
    # any other types of inverse
    ...
```
But I think this starts to get messy and is probably overkill for what we want. And no, we don't have different inverses depending on the network type.
Hi @rijobro ,

Could you please help share some comments about the TTA plan here?

Thanks in advance.

Hi @han-liu we certainly have it on our agenda to extend invertible transforms to include more than just spatial ones. 
Which intensity transform in particular are you using? Perhaps I can make this one invertible to be used as an example for future intensity transforms.
I guess invertible transforms don't have to be necessarilly mathematically invertible. We just need to define a suitable pseudo-inverse for it.

The example of adding Gaussian noise on the image intensities is interesting to consider. We could either make it an invertible transform by storing the noise added at test-time or just forget about the noise values and consider that the pseudo-inverse is the identity mapping. In any case, if TTA with Gaussian noise on the image intensities is to be used for segmentation purposes, the noise values are irrelevant (just a nuisance paraneters) for the segmentation output. In this case the pseudo-inverse to be applied to the network output is the identity.

Do we already handle different specialisation of the inverse transformation based on the type of network output?

In terms of really non-invertible transforms (those where storing the parameters is not enough to revert it up to numerical issues), I don't have a realistic example in mind. However, I guess that considering an image intensity change transform whereby we take the absolute value of the input intensity is a sufficient test case. In this scenario, we might want to have no defined inverse if the network is an auto-encoder but consider an identity pseudo-inverse if the network output is a segmentation map.

By the way, this raise the question of what to do for networks which are note image-in-image-out, e.g. classifiers. There as well, we might need to specialise the definition of inverse.

So in short, we should probably include some knowledge of the codomain and whether the transform acts of the input domain and output codomain in an [equivariant](https://en.wikipedia.org/wiki/Equivariant_map) fashion.
@rijobro @tvercaut Thanks for the clarification and I think what @tvercaut proposed here is a great feature to be added to TTA.
@tvercaut to take the example of `GaussianNoise`, the inverse of the spatial component of that transform is identity, because obviously adding Gaussian noise doesn't modify the image spatially. Having said that, I think it would be complicated for us to therefore say that a pseudo of the inverse can be given as identity. Because at that point, we're saying that the most important component of the inverse is the spatial component. This may or may not be true, but I'm not sure we really want to make a decision either way. 

A potential solution might be to have something like:

```python
def inverse(self, x):
    x = self.spatial_inverse(x)
    x = self.intensity_inverse(x)
    # any other types of inverse
    ...
```
But I think this starts to get messy and is probably overkill for what we want. And no, we don't have different inverses depending on the network type.

Patch:
diff --git a/monai/transforms/croppad/batch.py b/monai/transforms/croppad/batch.py
--- a/monai/transforms/croppad/batch.py
+++ b/monai/transforms/croppad/batch.py
@@ -20,10 +20,8 @@
 import torch
 
 from monai.data.utils import list_data_collate
-from monai.transforms.compose import Compose
 from monai.transforms.croppad.array import CenterSpatialCrop, SpatialPad
 from monai.transforms.inverse import InvertibleTransform
-from monai.transforms.utility.array import ToTensor
 from monai.utils.enums import InverseKeys, Method, NumpyPadMode
 
 __all__ = ["PadListDataCollate"]
@@ -95,20 +93,12 @@ def __call__(self, batch: Any):
             # If all same size, skip
             if np.all(np.array(max_shapes).min(axis=0) == max_shape):
                 continue
-            # Do we need to convert output to Tensor?
-            output_to_tensor = isinstance(batch[0][key_or_idx], torch.Tensor)
-
-            # Use `SpatialPadd` or `SpatialPad` to match sizes
-            # Default params are central padding, padding with 0's
-            # If input is dictionary, use the dictionary version so that the transformation is recorded
 
+            # Use `SpatialPad` to match sizes, Default params are central padding, padding with 0's
             padder = SpatialPad(spatial_size=max_shape, method=self.method, mode=self.mode, **self.np_kwargs)
-            transform = padder if not output_to_tensor else Compose([padder, ToTensor()])
-
             for idx, batch_i in enumerate(batch):
-                im = batch_i[key_or_idx]
-                orig_size = im.shape[1:]
-                padded = transform(batch_i[key_or_idx])
+                orig_size = batch_i[key_or_idx].shape[1:]
+                padded = padder(batch_i[key_or_idx])
                 batch = replace_element(padded, batch, idx, key_or_idx)
 
                 # If we have a dictionary of data, append to list
@@ -128,8 +118,10 @@ def inverse(data: dict) -> Dict[Hashable, np.ndarray]:
             transform_key = str(key) + InverseKeys.KEY_SUFFIX
             if transform_key in d:
                 transform = d[transform_key][-1]
-                if transform[InverseKeys.CLASS_NAME] == PadListDataCollate.__name__:
-                    d[key] = CenterSpatialCrop(transform["orig_size"])(d[key])
+                if not isinstance(transform, Dict):
+                    continue
+                if transform.get(InverseKeys.CLASS_NAME) == PadListDataCollate.__name__:
+                    d[key] = CenterSpatialCrop(transform.get("orig_size", -1))(d[key])  # fallback to image size
                     # remove transform
                     d[transform_key].pop()
         return d
diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py
--- a/monai/transforms/spatial/dictionary.py
+++ b/monai/transforms/spatial/dictionary.py
@@ -256,7 +256,7 @@ def __call__(
                     "old_affine": old_affine,
                     "mode": mode.value if isinstance(mode, Enum) else mode,
                     "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode,
-                    "align_corners": align_corners if align_corners is not None else "none",
+                    "align_corners": align_corners if align_corners is not None else InverseKeys.NONE,
                 },
                 orig_size=original_spatial_shape,
             )
@@ -288,7 +288,7 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd
                 affine=meta_data["affine"],  # type: ignore
                 mode=mode,
                 padding_mode=padding_mode,
-                align_corners=False if align_corners == "none" else align_corners,
+                align_corners=False if align_corners == InverseKeys.NONE else align_corners,
                 dtype=dtype,
                 output_spatial_shape=orig_size,
             )
@@ -562,7 +562,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
                 key,
                 extra_info={
                     "mode": mode.value if isinstance(mode, Enum) else mode,
-                    "align_corners": align_corners if align_corners is not None else "none",
+                    "align_corners": align_corners if align_corners is not None else InverseKeys.NONE,
                 },
             )
             d[key] = self.resizer(d[key], mode=mode, align_corners=align_corners)
@@ -577,7 +577,9 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd
             align_corners = transform[InverseKeys.EXTRA_INFO]["align_corners"]
             # Create inverse transform
             inverse_transform = Resize(
-                spatial_size=orig_size, mode=mode, align_corners=None if align_corners == "none" else align_corners
+                spatial_size=orig_size,
+                mode=mode,
+                align_corners=None if align_corners == InverseKeys.NONE else align_corners,
             )
             # Apply inverse transform
             d[key] = inverse_transform(d[key])
@@ -1346,7 +1348,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
                     "rot_mat": rot_mat,
                     "mode": mode.value if isinstance(mode, Enum) else mode,
                     "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode,
-                    "align_corners": align_corners if align_corners is not None else "none",
+                    "align_corners": align_corners if align_corners is not None else InverseKeys.NONE,
                 },
             )
         return d
@@ -1366,7 +1368,7 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd
                 normalized=False,
                 mode=mode,
                 padding_mode=padding_mode,
-                align_corners=False if align_corners == "none" else align_corners,
+                align_corners=False if align_corners == InverseKeys.NONE else align_corners,
                 reverse_indexing=True,
             )
             img_t: torch.Tensor
@@ -1478,7 +1480,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
                     "rot_mat": rot_mat,
                     "mode": mode.value if isinstance(mode, Enum) else mode,
                     "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode,
-                    "align_corners": align_corners if align_corners is not None else "none",
+                    "align_corners": align_corners if align_corners is not None else InverseKeys.NONE,
                 },
             )
         return d
@@ -1500,7 +1502,7 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd
                     normalized=False,
                     mode=mode,
                     padding_mode=padding_mode,
-                    align_corners=False if align_corners == "none" else align_corners,
+                    align_corners=False if align_corners == InverseKeys.NONE else align_corners,
                     reverse_indexing=True,
                 )
                 img_t: torch.Tensor
@@ -1578,7 +1580,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
                 extra_info={
                     "mode": mode.value if isinstance(mode, Enum) else mode,
                     "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode,
-                    "align_corners": align_corners if align_corners is not None else "none",
+                    "align_corners": align_corners if align_corners is not None else InverseKeys.NONE,
                 },
             )
             d[key] = self.zoomer(d[key], mode=mode, padding_mode=padding_mode, align_corners=align_corners)
@@ -1599,7 +1601,7 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd
                 d[key],
                 mode=mode,
                 padding_mode=padding_mode,
-                align_corners=None if align_corners == "none" else align_corners,
+                align_corners=None if align_corners == InverseKeys.NONE else align_corners,
             )
             # Size might be out by 1 voxel so pad
             d[key] = SpatialPad(transform[InverseKeys.ORIG_SIZE], mode="edge")(d[key])  # type: ignore
@@ -1697,7 +1699,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
                     "zoom": self.rand_zoom._zoom,
                     "mode": mode.value if isinstance(mode, Enum) else mode,
                     "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode,
-                    "align_corners": align_corners if align_corners is not None else "none",
+                    "align_corners": align_corners if align_corners is not None else InverseKeys.NONE,
                 },
             )
         return d
@@ -1719,7 +1721,7 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd
                     d[key],
                     mode=mode,
                     padding_mode=padding_mode,
-                    align_corners=None if align_corners == "none" else align_corners,
+                    align_corners=None if align_corners == InverseKeys.NONE else align_corners,
                 )
                 # Size might be out by 1 voxel so pad
                 d[key] = SpatialPad(transform[InverseKeys.ORIG_SIZE], mode="edge")(d[key])  # type: ignore
diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py
--- a/monai/transforms/utils.py
+++ b/monai/transforms/utils.py
@@ -1207,7 +1207,7 @@ def convert_inverse_interp_mode(trans_info: List, mode: str = "nearest", align_c
     interp_modes = [i.value for i in InterpolateMode] + [i.value for i in GridSampleMode]
 
     # set to string for DataLoader collation
-    align_corners_ = "none" if align_corners is None else align_corners
+    align_corners_ = InverseKeys.NONE if align_corners is None else align_corners
 
     for item in ensure_tuple(trans_info):
         if InverseKeys.EXTRA_INFO in item:
diff --git a/monai/utils/enums.py b/monai/utils/enums.py
--- a/monai/utils/enums.py
+++ b/monai/utils/enums.py
@@ -217,6 +217,7 @@ class InverseKeys:
     EXTRA_INFO = "extra_info"
     DO_TRANSFORM = "do_transforms"
     KEY_SUFFIX = "_transforms"
+    NONE = "none"
 
 
 class CommonKeys:


Test patch:
diff --git a/monai/data/test_time_augmentation.py b/monai/data/test_time_augmentation.py
--- a/monai/data/test_time_augmentation.py
+++ b/monai/data/test_time_augmentation.py
@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -36,6 +37,10 @@
 __all__ = ["TestTimeAugmentation"]
 
 
+def _identity(x):
+    return x
+
+
 class TestTimeAugmentation:
     """
     Class for performing test time augmentations. This will pass the same image through the network multiple times.
@@ -93,8 +98,8 @@ def __init__(
         self,
         transform: InvertibleTransform,
         batch_size: int,
-        num_workers: int,
-        inferrer_fn: Callable,
+        num_workers: int = 0,
+        inferrer_fn: Callable = _identity,
         device: Union[str, torch.device] = "cpu",
         image_key=CommonKeys.IMAGE,
         orig_key=CommonKeys.LABEL,
@@ -133,8 +138,8 @@ def _check_transforms(self):
         # check that whenever randoms is True, invertibles is also true
         for r, i in zip(randoms, invertibles):
             if r and not i:
-                raise RuntimeError(
-                    f"All applied random transform(s) must be invertible. Problematic transform: {type(r).__name__}"
+                warnings.warn(
+                    f"Not all applied random transform(s) are invertible. Problematic transform: {type(r).__name__}"
                 )
 
     def __call__(
@@ -161,7 +166,7 @@ def __call__(
         # generate batch of data of size == batch_size, dataset and dataloader
         data_in = [deepcopy(d) for _ in range(num_examples)]
         ds = Dataset(data_in, self.transform)
-        dl = DataLoader(ds, self.num_workers, batch_size=self.batch_size, collate_fn=pad_list_data_collate)
+        dl = DataLoader(ds, num_workers=self.num_workers, batch_size=self.batch_size, collate_fn=pad_list_data_collate)
 
         transform_key = self.orig_key + InverseKeys.KEY_SUFFIX
 
@@ -180,8 +185,10 @@ def __call__(
                 batch_output = batch_output.detach().cpu()
             if isinstance(batch_output, np.ndarray):
                 batch_output = torch.Tensor(batch_output)
-
-            transform_info = batch_data[transform_key]
+            transform_info = batch_data.get(transform_key, None)
+            if transform_info is None:
+                # no invertible transforms, adding dummy info for identity invertible
+                transform_info = [[InverseKeys.NONE] for _ in range(self.batch_size)]
             if self.nearest_interp:
                 transform_info = convert_inverse_interp_mode(
                     trans_info=deepcopy(transform_info), mode="nearest", align_corners=None
diff --git a/tests/test_testtimeaugmentation.py b/tests/test_testtimeaugmentation.py
--- a/tests/test_testtimeaugmentation.py
+++ b/tests/test_testtimeaugmentation.py
@@ -21,9 +21,18 @@
 from monai.data.utils import pad_list_data_collate
 from monai.losses import DiceLoss
 from monai.networks.nets import UNet
-from monai.transforms import Activations, AddChanneld, AsDiscrete, Compose, CropForegroundd, DivisiblePadd, RandAffined
+from monai.transforms import (
+    Activations,
+    AddChanneld,
+    AsDiscrete,
+    Compose,
+    CropForegroundd,
+    DivisiblePadd,
+    RandAffined,
+    RandScaleIntensityd,
+)
 from monai.transforms.croppad.dictionary import SpatialPadd
-from monai.transforms.spatial.dictionary import Rand2DElasticd, RandFlipd, Spacingd
+from monai.transforms.spatial.dictionary import RandFlipd, Spacingd
 from monai.utils import optional_import, set_determinism
 from tests.utils import TEST_NDARRAYS
 
@@ -46,13 +55,13 @@ def get_data(num_examples, input_size, data_type=np.asarray, include_label=True)
             create_test_image_2d, *input_size, rad_max=7, num_seg_classes=1, num_objs=1
         )
         data = []
-        for _ in range(num_examples):
+        for i in range(num_examples):
             im, label = custom_create_test_image_2d()
             d = {}
-            d["image"] = data_type(im)
+            d["image"] = data_type(im[:, i:])
             d["image_meta_dict"] = {"affine": np.eye(4)}
             if include_label:
-                d["label"] = data_type(label)
+                d["label"] = data_type(label[:, i:])
                 d["label_meta_dict"] = {"affine": np.eye(4)}
             data.append(d)
         return data[0] if num_examples == 1 else data
@@ -132,10 +141,20 @@ def test_fail_non_random(self):
         with self.assertRaises(RuntimeError):
             TestTimeAugmentation(transforms, None, None, None)
 
-    def test_fail_random_but_not_invertible(self):
-        transforms = Compose([AddChanneld("im"), Rand2DElasticd("im", None, None)])
-        with self.assertRaises(RuntimeError):
-            TestTimeAugmentation(transforms, None, None, None)
+    def test_warn_random_but_has_no_invertible(self):
+        transforms = Compose(
+            [AddChanneld("image"), RandFlipd("image", prob=1.0), RandScaleIntensityd("image", 0.1, prob=1.0)]
+        )
+        with self.assertWarns(UserWarning):
+            tta = TestTimeAugmentation(transforms, 5, 0, orig_key="image")
+            tta(self.get_data(1, (20, 20), data_type=np.float32))
+
+    def test_warn_random_but_all_not_invertible(self):
+        """test with no invertible stack"""
+        transforms = Compose([AddChanneld("image"), RandScaleIntensityd("image", 0.1, prob=1.0)])
+        with self.assertWarns(UserWarning):
+            tta = TestTimeAugmentation(transforms, 1, 0, orig_key="image")
+            tta(self.get_data(1, (20, 20), data_type=np.float32))
 
     def test_single_transform(self):
         for p in TEST_NDARRAYS: