Repository: Project-MONAI/MONAI

Issue:
PadListDataCollate to support inverse MetaTensor
`PadListDataCollate.inverse` should support MetaTensor (as it's used in test time augmentation)
https://github.com/Project-MONAI/MONAI/blob/99b9dd7cf0667113e56f0dd90d40699115048d59/monai/transforms/croppad/batch.py#L114-L118

also the test cases are incomplete:
https://github.com/Project-MONAI/MONAI/blob/99b9dd7cf0667113e56f0dd90d40699115048d59/tests/test_pad_collation.py#L97-L106
PadListDataCollate to support inverse MetaTensor
`PadListDataCollate.inverse` should support MetaTensor (as it's used in test time augmentation)
https://github.com/Project-MONAI/MONAI/blob/99b9dd7cf0667113e56f0dd90d40699115048d59/monai/transforms/croppad/batch.py#L114-L118

also the test cases are incomplete:
https://github.com/Project-MONAI/MONAI/blob/99b9dd7cf0667113e56f0dd90d40699115048d59/tests/test_pad_collation.py#L97-L106


Discussion:

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
@@ -19,6 +19,7 @@
 import numpy as np
 import torch
 
+from monai.data.meta_tensor import MetaTensor
 from monai.data.utils import list_data_collate
 from monai.transforms.croppad.array import CenterSpatialCrop, SpatialPad
 from monai.transforms.inverse import InvertibleTransform
@@ -112,13 +113,18 @@ def inverse(data: dict) -> Dict[Hashable, np.ndarray]:
 
         d = deepcopy(data)
         for key in d:
-            transform_key = InvertibleTransform.trace_key(key)
-            if transform_key in d:
-                transform = d[transform_key][-1]
-                if not isinstance(transform, Dict):
-                    continue
-                if transform.get(TraceKeys.CLASS_NAME) == PadListDataCollate.__name__:
-                    d[key] = CenterSpatialCrop(transform.get("orig_size", -1))(d[key])  # fallback to image size
-                    # remove transform
-                    d[transform_key].pop()
+            transforms = None
+            if isinstance(d[key], MetaTensor):
+                transforms = d[key].applied_operations
+            else:
+                transform_key = InvertibleTransform.trace_key(key)
+                if transform_key in d:
+                    transforms = d[transform_key]
+            if not transforms or not isinstance(transforms[-1], Dict):
+                continue
+            if transforms[-1].get(TraceKeys.CLASS_NAME) == PadListDataCollate.__name__:
+                xform = transforms.pop()
+                cropping = CenterSpatialCrop(xform.get(TraceKeys.ORIG_SIZE, -1))
+                with cropping.trace_transform(False):
+                    d[key] = cropping(d[key])  # fallback to image size
         return d


Test patch:
diff --git a/tests/test_pad_collation.py b/tests/test_pad_collation.py
--- a/tests/test_pad_collation.py
+++ b/tests/test_pad_collation.py
@@ -98,9 +98,12 @@ def test_pad_collation(self, t_type, collate_method, transform):
         # check collation in forward direction
         for data in loader:
             if t_type == dict:
+                shapes = []
                 decollated_data = decollate_batch(data)
                 for d in decollated_data:
-                    PadListDataCollate.inverse(d)
+                    output = PadListDataCollate.inverse(d)
+                    shapes.append(output["image"].shape)
+                self.assertTrue(len(set(shapes)) > 1)  # inverted shapes must be different because of random xforms
 
 
 if __name__ == "__main__":
diff --git a/tests/test_testtimeaugmentation.py b/tests/test_testtimeaugmentation.py
--- a/tests/test_testtimeaugmentation.py
+++ b/tests/test_testtimeaugmentation.py
@@ -58,9 +58,7 @@ def get_data(num_examples, input_size, data_type=np.asarray, include_label=True)
         data = []
         for i in range(num_examples):
             im, label = custom_create_test_image_2d()
-            d = {}
-            d["image"] = data_type(im[:, i:])
-            d[PostFix.meta("image")] = {"affine": np.eye(4)}
+            d = {"image": data_type(im[:, i:])}
             if include_label:
                 d["label"] = data_type(label[:, i:])
                 d[PostFix.meta("label")] = {"affine": np.eye(4)}