Repository: Project-MONAI/MONAI

Issue:
TransformInverter has `collate` arg but can't work
**Describe the bug**
Currently, the `TransformInverter` hander can't support collate that doesn't generate 1 list of data.



Discussion:

Patch:
diff --git a/monai/handlers/iteration_metric.py b/monai/handlers/iteration_metric.py
--- a/monai/handlers/iteration_metric.py
+++ b/monai/handlers/iteration_metric.py
@@ -31,6 +31,8 @@ class IterationMetric(Metric):  # type: ignore[valid-type, misc] # due to option
     Class for metrics that should be computed on every iteration and compute final results when epoch completed.
     Similar to the `EpochMetric` in ignite:
     https://github.com/pytorch/ignite/blob/v0.4.2/ignite/metrics/epoch_metric.py#L13.
+    The input `prediction` or `label` data can be a PyTorch Tensor or numpy array with batch dim and channel dim,
+    or a list of PyTorch Tensor or numpy array without batch dim.
 
     Args:
         metric_fn: callable function or class to compute raw metric results after every iteration.
diff --git a/monai/handlers/segmentation_saver.py b/monai/handlers/segmentation_saver.py
--- a/monai/handlers/segmentation_saver.py
+++ b/monai/handlers/segmentation_saver.py
@@ -30,6 +30,7 @@ class SegmentationSaver:
     Event handler triggered on completing every iteration to save the segmentation predictions into files.
     It can extract the input image meta data(filename, affine, original_shape, etc.) and resample the predictions
     based on the meta data.
+    The predictions can be PyTorch Tensor with [B, C, H, W, [D]] shape or a list of Tensor without batch dim.
 
     """
 
@@ -117,6 +118,7 @@ def __init__(
             scale=scale,
             dtype=dtype,
             output_dtype=output_dtype,
+            save_batch=True,
             squeeze_end_dims=squeeze_end_dims,
             data_root_dir=data_root_dir,
         )
@@ -146,19 +148,5 @@ def __call__(self, engine: Engine) -> None:
         """
         meta_data = self.batch_transform(engine.state.batch)
         engine_output = self.output_transform(engine.state.output)
-        if isinstance(engine_output, (tuple, list)):
-            # if a list of data in shape: [channel, H, W, [D]], save every item separately
-            self._saver.save_batch = False
-            for i, d in enumerate(engine_output):
-                if isinstance(meta_data, dict):
-                    meta_ = {k: meta_data[k][i] for k in meta_data}
-                elif isinstance(meta_data, (list, tuple)):
-                    meta_ = meta_data[i]
-                else:
-                    meta_ = meta_data
-                self._saver(d, meta_)
-        else:
-            # if the data is in shape: [batch, channel, H, W, [D]]
-            self._saver.save_batch = True
-            self._saver(engine_output, meta_data)
+        self._saver(engine_output, meta_data)
         self.logger.info("saved all the model outputs into files.")
diff --git a/monai/handlers/transform_inverter.py b/monai/handlers/transform_inverter.py
--- a/monai/handlers/transform_inverter.py
+++ b/monai/handlers/transform_inverter.py
@@ -65,8 +65,8 @@ def __init__(
                 default is `meta_dict`, the meta data is a dictionary object.
                 For example, to handle key `image`,  read/write affine matrices from the
                 metadata `image_meta_dict` dictionary's `affine` field.
-            collate_fn: how to collate data after inverse transformations.
-                default won't do any collation, so the output will be a list of size batch size.
+            collate_fn: how to collate data after inverse transformations. default won't do any collation,
+                so the output will be a list of PyTorch Tensor or numpy array without batch dim.
             postfix: will save the inverted result into `ignite.engine.output` with key `{output_key}_{postfix}`.
             nearest_interp: whether to use `nearest` interpolation mode when inverting the spatial transforms,
                 default to `True`. If `False`, use the same interpolation mode as the original transform.
diff --git a/monai/transforms/io/array.py b/monai/transforms/io/array.py
--- a/monai/transforms/io/array.py
+++ b/monai/transforms/io/array.py
@@ -169,8 +169,9 @@ class SaveImage(Transform):
     Save transformed data into files, support NIfTI and PNG formats.
     It can work for both numpy array and PyTorch Tensor in both pre-transform chain
     and post transform chain.
+    It can also save a list of PyTorch Tensor or numpy array without `batch dim`.
 
-    NB: image should include channel dimension: [B],C,H,W,[D].
+    Note: image should include channel dimension: [B],C,H,W,[D].
 
     Args:
         output_dir: output image directory.
@@ -276,7 +277,19 @@ def __call__(self, img: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dic
             meta_data: key-value pairs of meta_data corresponding to the data.
 
         """
-        if self.save_batch:
-            self.saver.save_batch(img, meta_data)
+        if isinstance(img, (tuple, list)):
+            # if a list of data in shape: [channel, H, W, [D]], save every item separately
+            meta_: Optional[Dict] = None
+            for i, d in enumerate(img):
+                if isinstance(meta_data, dict):
+                    meta_ = {k: meta_data[k][i] for k in meta_data}
+                elif isinstance(meta_data, (list, tuple)):
+                    meta_ = meta_data[i]
+                else:
+                    meta_ = meta_data
+                self.saver.save(d, meta_)
         else:
-            self.saver.save(img, meta_data)
+            if self.save_batch:
+                self.saver.save_batch(img, meta_data)
+            else:
+                self.saver.save(img, meta_data)
diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py
--- a/monai/transforms/post/dictionary.py
+++ b/monai/transforms/post/dictionary.py
@@ -419,6 +419,10 @@ class Invertd(MapTransform):
     A typical usage is to invert the pre-transforms (appplied on input `image`) on the model `pred` data.
 
     Note:
+        According to the `collate_fn`, this transform may return a list of Tensor without batch dim,
+        thus some following post transforms may not support a list of Tensor, and users can leverage the
+        `post_func` arg for basic processing logic.
+
         As this transform only accepts 1 input dict while ignite stores model input data in `state.batch`
         and stores model output data in `state.output`, so it's not compatible with MONAI engines so far.
         For MONAI workflow engines, please use the `TransformInverter` handler instead.
@@ -455,8 +459,8 @@ def __init__(
                 default is `meta_dict`, the meta data is a dictionary object.
                 For example, to handle orig_key `image`,  read/write `affine` matrices from the
                 metadata `image_meta_dict` dictionary's `affine` field.
-            collate_fn: how to collate data after inverse transformations.
-                default won't do any collation, so the output will be a list of size batch size.
+            collate_fn: how to collate data after inverse transformations. default won't do any collation,
+                so the output will be a list of PyTorch Tensor or numpy array without batch dim.
             postfix: will save the inverted result into dict with key `{key}_{postfix}`.
             nearest_interp: whether to use `nearest` interpolation mode when inverting the spatial transforms,
                 default to `True`. If `False`, use the same interpolation mode as the original transform.
@@ -526,13 +530,20 @@ def __call__(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]:
 
             # save the inverted data
             inverted_key = f"{key}_{self.postfix}"
-            d[inverted_key] = [
-                post_func(self._totensor(i[orig_key]).to(device) if to_tensor else i[orig_key]) for i in inverted
-            ]
-
-            # save the inverted meta dict
-            if meta_dict_key in d:
-                d[f"{inverted_key}_{self.meta_key_postfix}"] = [i.get(meta_dict_key) for i in inverted]
+            if isinstance(inverted, (tuple, list)):
+                d[inverted_key] = [
+                    post_func(self._totensor(i[orig_key]).to(device) if to_tensor else i[orig_key]) for i in inverted
+                ]
+                # save the inverted meta dict
+                if meta_dict_key in d:
+                    d[f"{inverted_key}_{self.meta_key_postfix}"] = [i.get(meta_dict_key) for i in inverted]
+            else:
+                d[inverted_key] = post_func(
+                    self._totensor(inverted[orig_key]).to(device) if to_tensor else inverted[orig_key]
+                )
+                # save the inverted meta dict
+                if meta_dict_key in d:
+                    d[f"{inverted_key}_{self.meta_key_postfix}"] = inverted.get(meta_dict_key)
         return d
 
 


Test patch:
diff --git a/tests/test_handler_transform_inverter.py b/tests/test_handler_transform_inverter.py
--- a/tests/test_handler_transform_inverter.py
+++ b/tests/test_handler_transform_inverter.py
@@ -16,7 +16,7 @@
 import torch
 from ignite.engine import Engine
 
-from monai.data import CacheDataset, DataLoader, create_test_image_3d
+from monai.data import CacheDataset, DataLoader, create_test_image_3d, pad_list_data_collate
 from monai.engines.utils import IterationEvents
 from monai.handlers import TransformInverter
 from monai.transforms import (
@@ -104,6 +104,7 @@ def _train_func(engine, batch):
             nearest_interp=[True, False],
             post_func=[lambda x: x + 10, lambda x: x],
             postfix="inverted2",
+            collate_fn=pad_list_data_collate,
             num_workers=0 if sys.platform == "darwin" or torch.cuda.is_available() else 2,
         ).attach(engine)
 
@@ -133,15 +134,15 @@ def _train_func(engine, batch):
         self.assertTrue((reverted.size - n_good) in (25300, 1812, 1824), "diff. in 3 possible values")
 
         # check the case that different items use different interpolation mode to invert transforms
-        for i in engine.state.output["image_inverted2"]:
-            # if the interpolation mode is nearest, accumulated diff should be smaller than 1
-            self.assertLess(torch.sum(i.to(torch.float) - i.to(torch.uint8).to(torch.float)).item(), 1.0)
-            self.assertTupleEqual(i.shape, (1, 100, 101, 107))
-
-        for i in engine.state.output["label_inverted2"]:
-            # if the interpolation mode is not nearest, accumulated diff should be greater than 10000
-            self.assertGreater(torch.sum(i.to(torch.float) - i.to(torch.uint8).to(torch.float)).item(), 10000.0)
-            self.assertTupleEqual(i.shape, (1, 100, 101, 107))
+        d = engine.state.output["image_inverted2"]
+        # if the interpolation mode is nearest, accumulated diff should be smaller than 1
+        self.assertLess(torch.sum(d.to(torch.float) - d.to(torch.uint8).to(torch.float)).item(), 1.0)
+        self.assertTupleEqual(d.shape, (2, 1, 100, 101, 107))
+
+        d = engine.state.output["label_inverted2"]
+        # if the interpolation mode is not nearest, accumulated diff should be greater than 10000
+        self.assertGreater(torch.sum(d.to(torch.float) - d.to(torch.uint8).to(torch.float)).item(), 10000.0)
+        self.assertTupleEqual(d.shape, (2, 1, 100, 101, 107))
 
 
 if __name__ == "__main__":
diff --git a/tests/test_save_image.py b/tests/test_save_image.py
--- a/tests/test_save_image.py
+++ b/tests/test_save_image.py
@@ -83,9 +83,37 @@
     False,
 ]
 
+TEST_CASE_7 = [
+    [torch.randint(0, 255, (1, 2, 3, 4)), torch.randint(0, 255, (1, 2, 3, 4))],
+    [{"filename_or_obj": "testfile0.nii.gz"}, {"filename_or_obj": "testfile1.nii.gz"}],
+    ".nii.gz",
+    False,
+    False,
+]
+
+TEST_CASE_8 = [
+    [torch.randint(0, 255, (1, 2, 3, 4))],
+    {"filename_or_obj": ["testfile0.nii.gz"]},
+    ".nii.gz",
+    False,
+    False,
+]
+
 
 class TestSaveImage(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6])
+    @parameterized.expand(
+        [
+            TEST_CASE_0,
+            TEST_CASE_1,
+            TEST_CASE_2,
+            TEST_CASE_3,
+            TEST_CASE_4,
+            TEST_CASE_5,
+            TEST_CASE_6,
+            TEST_CASE_7,
+            TEST_CASE_8,
+        ]
+    )
     def test_saved_content(self, test_data, meta_data, output_ext, resample, save_batch):
         with tempfile.TemporaryDirectory() as tempdir:
             trans = SaveImage(