Repository: Project-MONAI/MONAI

Issue:
Change `Invertd` transform to be `in-place` operation
**Is your feature request related to a problem? Please describe.**
Currently, all the MONAI transforms are in place operation except the newly added `Invertd` transform.
Usually, if want to keep the previous value, users will add `CopyItemsd` transform first.
So I would suggest to remove the `postfix` arg of `Invertd` transform to make it as `in-place` operation and align with other transforms.



Discussion:

Resolving patch:
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
@@ -32,9 +32,9 @@ class TransformInverter:
     Ignite handler to automatically invert `transforms`.
     It takes `engine.state.output` as the input data and uses the transforms information from `engine.state.batch`.
     Expect both `engine.state.output` and `engine.state.batch` to be dictionary data.
-    The inverted data are stored in `engine.state.output` with key: "{output_key}_{postfix}".
+    The inverted data is in-place saved back to `engine.state.output` with key: "{output_key}".
     And the inverted meta dict will be stored in `engine.state.batch`
-    with key: "{output_key}_{postfix}_{meta_key_postfix}".
+    with key: "{meta_keys}" or "{key}_{meta_key_postfix}".
 
     """
 
@@ -45,9 +45,9 @@ def __init__(
         output_keys: KeysCollection = CommonKeys.PRED,
         batch_keys: KeysCollection = CommonKeys.IMAGE,
         meta_keys: Optional[KeysCollection] = None,
+        batch_meta_keys: Optional[KeysCollection] = None,
         meta_key_postfix: str = "meta_dict",
         collate_fn: Optional[Callable] = no_collation,
-        postfix: str = "inverted",
         nearest_interp: Union[bool, Sequence[bool]] = True,
         to_tensor: Union[bool, Sequence[bool]] = True,
         device: Union[Union[str, torch.device], Sequence[Union[str, torch.device]]] = "cpu",
@@ -59,24 +59,29 @@ def __init__(
             transform: a callable data transform on input data.
             loader: data loader used to run transforms and generate the batch of data.
             output_keys: the key of expected data in `ignite.engine.output`, invert transforms on it.
-                it also can be a list of keys, will invert transform for each of them. Default to "pred".
+                it also can be a list of keys, will invert transform for each of them.
+                Default to "pred". it's in-place operation.
             batch_keys: the key of input data in `ignite.engine.batch`. will get the applied transforms
                 for this input data, then invert them for the expected data with `output_keys`.
                 It can also be a list of keys, each matches to the `output_keys` data. default to "image".
-            meta_keys: explicitly indicate the key of the corresponding meta data dictionary.
-                for example, for data with key `image`, the metadata by default is in `image_meta_dict`.
+            meta_keys: explicitly indicate the key for the inverted meta data dictionary.
                 the meta data is a dictionary object which contains: filename, original_shape, etc.
                 it can be a sequence of string, map to the `keys`.
-                if None, will try to construct meta_keys by `key_{meta_key_postfix}`.
-                it also can be a list of string, each matches to the `output_keys` data.
-            meta_key_postfix: if meta_keys is None, use `{batch_key}_{postfix}` to to fetch the meta data according
-                to the key data, default is `meta_dict`, the meta data is a dictionary object.
-                For example, to handle key `image`,  read/write affine matrices from the
+                if None, will try to construct meta_keys by `{key}_{meta_key_postfix}`.
+            batch_meta_keys: the key of the meta data of input data in `ignite.engine.batch`,
+                will get the `affine`, `data_shape`, etc.
+                the meta data is a dictionary object which contains: filename, original_shape, etc.
+                it can be a sequence of string, map to the `keys`.
+                if None, will try to construct meta_keys by `{orig_key}_{meta_key_postfix}`.
+                meta data will also be inverted and stored in `meta_keys`.
+            meta_key_postfix: if `orig_meta_keys` is None, use `{orig_key}_{meta_key_postfix}` to to fetch the
+                meta data from dict, if `meta_keys` is None, use `{key}_{meta_key_postfix}`.
+                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.
-                the inverted meta dict will be stored with key: "{output_key}_{postfix}_{meta_key_postfix}".
+                the inverted meta dict will be stored with key: "{key}_{meta_key_postfix}".
             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.
                 it also can be a list of bool, each matches to the `output_keys` data.
@@ -98,9 +103,9 @@ def __init__(
             loader=loader,
             orig_keys=batch_keys,
             meta_keys=meta_keys,
+            orig_meta_keys=batch_meta_keys,
             meta_key_postfix=meta_key_postfix,
             collate_fn=collate_fn,
-            postfix=postfix,
             nearest_interp=nearest_interp,
             to_tensor=to_tensor,
             device=device,
@@ -108,8 +113,10 @@ def __init__(
             num_workers=num_workers,
         )
         self.output_keys = ensure_tuple(output_keys)
+        self.meta_keys = ensure_tuple_rep(None, len(self.output_keys)) if meta_keys is None else ensure_tuple(meta_keys)
+        if len(self.output_keys) != len(self.meta_keys):
+            raise ValueError("meta_keys should have the same length as output_keys.")
         self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.output_keys))
-        self.postfix = postfix
 
     def attach(self, engine: Engine) -> None:
         """
@@ -128,12 +135,12 @@ def __call__(self, engine: Engine) -> None:
         data.update(engine.state.output)
         ret = self.inverter(data)
 
-        for output_key, meta_key_postfix in zip(self.output_keys, self.meta_key_postfix):
+        for output_key, meta_key, meta_key_postfix in zip(self.output_keys, self.meta_keys, self.meta_key_postfix):
             # save the inverted data into state.output
-            inverted_key = f"{output_key}_{self.postfix}"
-            if inverted_key in ret:
-                engine.state.output[inverted_key] = ret[inverted_key]
+            engine.state.output[output_key] = ret.get(output_key)
             # save the inverted meta dict into state.batch
-            meta_dict_key = f"{inverted_key}_{meta_key_postfix}"
-            if meta_dict_key in ret:
-                engine.state.batch[meta_dict_key] = ret[meta_dict_key]
+            meta_key = meta_key or f"{output_key}_{meta_key_postfix}"
+            if meta_key in ret:
+                # FIXME: we save inverted meta dict into `batch` to be compatible with `SegmentationSaver`
+                # will deprecate both handlers soon
+                engine.state.batch[meta_key] = ret.get(meta_key)
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
@@ -40,7 +40,7 @@
 from monai.transforms.transform import MapTransform
 from monai.transforms.utility.array import ToTensor
 from monai.transforms.utils import allow_missing_keys_mode, convert_inverse_interp_mode
-from monai.utils import ensure_tuple_rep
+from monai.utils import ensure_tuple, ensure_tuple_rep
 from monai.utils.enums import InverseKeys
 
 __all__ = [
@@ -416,21 +416,22 @@ def __call__(self, data: Mapping[Hashable, Union[np.ndarray, torch.Tensor]]):
 
 class Invertd(MapTransform):
     """
-    Utility transform to automatically invert the previous transforms based on transform information.
-    It extracts the transform information applied on the data with `orig_keys`, then inverts the transforms
-    on the data with `keys`. several `keys` can share one `orig_keys`.
-    A typical usage is to invert the pre-transforms (appplied on input `image`) on the model `pred` data.
+    Utility transform to automatically invert the previously applied transforms.
+
+    The output of the inverted data and metadata will be stored at `keys` and `meta_keys` respectively.
+    To correctly invert the transforms, the information of the previously applied transforms should be
+    available at `orig_keys`, and the original metadata at `orig_meta_keys`.
+    (`meta_key_postfix` is an optional string to conveniently construct "meta_keys" and/or "orig_meta_keys".)
+
+    Typical usage is to invert the pre-transforms (applied on input `image`) on the model `pred` data.
+    A detailed usage example is available in the tutorial:
+    https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/torch/unet_inference_dict.py
 
     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.
-        Users can use this transform in a regular PyTorch program which uses dict data for transforms.
-
     """
 
     def __init__(
@@ -440,9 +441,9 @@ def __init__(
         loader: TorchDataLoader,
         orig_keys: KeysCollection,
         meta_keys: Optional[KeysCollection] = None,
+        orig_meta_keys: Optional[KeysCollection] = None,
         meta_key_postfix: str = "meta_dict",
         collate_fn: Optional[Callable] = no_collation,
-        postfix: str = "inverted",
         nearest_interp: Union[bool, Sequence[bool]] = True,
         to_tensor: Union[bool, Sequence[bool]] = True,
         device: Union[Union[str, torch.device], Sequence[Union[str, torch.device]]] = "cpu",
@@ -452,26 +453,30 @@ def __init__(
     ) -> None:
         """
         Args:
-            keys: the key of expected data in the dict, invert transforms on it.
+            keys: the key of expected data in the dict, invert transforms on it, in-place operation.
                 it also can be a list of keys, will invert transform for each of them, like: ["pred", "pred_class2"].
             transform: the previous callable transform that applied on input data.
             loader: data loader used to run transforms and generate the batch of data.
             orig_keys: the key of the original input data in the dict. will get the applied transform information
                 for this input data, then invert them for the expected data with `keys`.
                 It can also be a list of keys, each matches to the `keys` data.
-            meta_keys: explicitly indicate the key of the corresponding meta data dictionary.
-                for example, for data with key `image`, the metadata by default is in `image_meta_dict`.
+            meta_keys: explicitly indicate the key for the inverted meta data dictionary.
                 the meta data is a dictionary object which contains: filename, original_shape, etc.
                 it can be a sequence of string, map to the `keys`.
-                if None, will try to construct meta_keys by `key_{meta_key_postfix}`.
-            meta_key_postfix: if meta_keys is None, use `{orig_key}_{postfix}` to to fetch the meta data from dict,
+                if None, will try to construct meta_keys by `{key}_{meta_key_postfix}`.
+            orig_meta_keys: the key of the meta data of original input data, will get the `affine`, `data_shape`, etc.
+                the meta data is a dictionary object which contains: filename, original_shape, etc.
+                it can be a sequence of string, map to the `keys`.
+                if None, will try to construct meta_keys by `{orig_key}_{meta_key_postfix}`.
+                meta data will also be inverted and stored in `meta_keys`.
+            meta_key_postfix: if `orig_meta_keys` is None, use `{orig_key}_{meta_key_postfix}` to to fetch the
+                meta data from dict, if `meta_keys` is None, use `{key}_{meta_key_postfix}`.
                 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.
-                the inverted meta dict will be stored with key: "{key}_{postfix}_{meta_key_postfix}".
+                the inverted meta dict will be stored with key: "{key}_{meta_key_postfix}".
             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.
                 it also can be a list of bool, each matches to the `keys` data.
@@ -497,9 +502,11 @@ def __init__(
             num_workers=num_workers,
         )
         self.orig_keys = ensure_tuple_rep(orig_keys, len(self.keys))
-        self.meta_keys = ensure_tuple_rep(meta_keys, len(self.keys))
+        self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys)
+        if len(self.keys) != len(self.meta_keys):
+            raise ValueError("meta_keys should have the same length as keys.")
+        self.orig_meta_keys = ensure_tuple_rep(orig_meta_keys, len(self.keys))
         self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys))
-        self.postfix = postfix
         self.nearest_interp = ensure_tuple_rep(nearest_interp, len(self.keys))
         self.to_tensor = ensure_tuple_rep(to_tensor, len(self.keys))
         self.device = ensure_tuple_rep(device, len(self.keys))
@@ -512,6 +519,7 @@ def __call__(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]:
             key,
             orig_key,
             meta_key,
+            orig_meta_key,
             meta_key_postfix,
             nearest_interp,
             to_tensor,
@@ -521,6 +529,7 @@ def __call__(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]:
             d,
             self.orig_keys,
             self.meta_keys,
+            self.orig_meta_keys,
             self.meta_key_postfix,
             self.nearest_interp,
             self.to_tensor,
@@ -548,29 +557,27 @@ def __call__(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]:
                 orig_key: input,
                 transform_key: transform_info,
             }
-            meta_dict_key = meta_key or f"{orig_key}_{meta_key_postfix}"
-            if meta_dict_key in d:
-                input_dict[meta_dict_key] = d[meta_dict_key]
+            orig_meta_key = orig_meta_key or f"{orig_key}_{meta_key_postfix}"
+            meta_key = meta_key or f"{key}_{meta_key_postfix}"
+            if orig_meta_key in d:
+                input_dict[orig_meta_key] = d[orig_meta_key]
 
             with allow_missing_keys_mode(self.transform):  # type: ignore
                 inverted = self.inverter(input_dict)
 
             # save the inverted data
-            inverted_key = f"{key}_{self.postfix}"
             if isinstance(inverted, (tuple, list)):
-                d[inverted_key] = [
+                d[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}_{meta_key_postfix}"] = [i.get(meta_dict_key) for i in inverted]
+                if orig_meta_key in d:
+                    d[meta_key] = [i.get(orig_meta_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]
-                )
+                d[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}_{meta_key_postfix}"] = inverted.get(meta_dict_key)
+                if orig_meta_key in d:
+                    d[meta_key] = inverted.get(orig_meta_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
@@ -23,6 +23,7 @@
     AddChanneld,
     CastToTyped,
     Compose,
+    CopyItemsd,
     LoadImaged,
     Orientationd,
     RandAffined,
@@ -62,6 +63,8 @@ def test_invert(self):
                 ResizeWithPadOrCropd(KEYS, 100),
                 ToTensord("image"),  # test to support both Tensor and Numpy array when inverting
                 CastToTyped(KEYS, dtype=[torch.uint8, np.uint8]),
+                CopyItemsd("label", times=2, names=["label_inverted1", "label_inverted2"]),
+                CopyItemsd("image", times=2, names=["image_inverted1", "image_inverted2"]),
             ]
         )
         data = [{"image": im_fname, "label": seg_fname} for _ in range(12)]
@@ -86,11 +89,11 @@ def _train_func(engine, batch):
         TransformInverter(
             transform=transform,
             loader=loader,
-            output_keys=["image", "label"],
+            output_keys=["image_inverted1", "label_inverted1"],
             batch_keys="label",
-            meta_keys="label_meta_dict",
+            meta_keys=["image_inverted1_meta_dict", "label_inverted1_meta_dict"],
+            batch_meta_keys="label_meta_dict",
             nearest_interp=True,
-            postfix="inverted1",
             to_tensor=[True, False],
             device="cpu",
             num_workers=0 if sys.platform == "darwin" or torch.cuda.is_available() else 2,
@@ -100,11 +103,13 @@ def _train_func(engine, batch):
         TransformInverter(
             transform=transform,
             loader=loader,
-            output_keys=["image", "label"],
+            output_keys=["image_inverted2", "label_inverted2"],
             batch_keys="image",
+            meta_keys=None,
+            batch_meta_keys="image_meta_dict",
+            meta_key_postfix="meta_dict",
             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)
@@ -125,7 +130,7 @@ def _train_func(engine, batch):
         reverted = engine.state.output["label_inverted1"][-1].astype(np.int32)
         original = LoadImaged(KEYS)(data[-1])["label"]
         n_good = np.sum(np.isclose(reverted, original, atol=1e-3))
-        reverted_name = engine.state.output["label_meta_dict"]["filename_or_obj"][-1]
+        reverted_name = engine.state.batch["label_inverted1_meta_dict"][-1]["filename_or_obj"]
         original_name = data[-1]["label"]
         self.assertEqual(reverted_name, original_name)
         print("invert diff", reverted.size - n_good)
diff --git a/tests/test_invertd.py b/tests/test_invertd.py
--- a/tests/test_invertd.py
+++ b/tests/test_invertd.py
@@ -20,6 +20,7 @@
     AddChanneld,
     CastToTyped,
     Compose,
+    CopyItemsd,
     Invertd,
     LoadImaged,
     Orientationd,
@@ -60,6 +61,7 @@ def test_invert(self):
                 ResizeWithPadOrCropd(KEYS, 100),
                 ToTensord("image"),  # test to support both Tensor and Numpy array when inverting
                 CastToTyped(KEYS, dtype=[torch.uint8, np.uint8]),
+                CopyItemsd("label", times=1, names="label_inverted"),
             ]
         )
         data = [{"image": im_fname, "label": seg_fname} for _ in range(12)]
@@ -70,13 +72,14 @@ def test_invert(self):
         dataset = CacheDataset(data, transform=transform, progress=False)
         loader = DataLoader(dataset, num_workers=num_workers, batch_size=5)
         inverter = Invertd(
-            keys=["image", "label"],
+            # `image` was not copied, invert the original value directly
+            keys=["image", "label_inverted"],
             transform=transform,
             loader=loader,
             orig_keys="label",
-            meta_keys="label_meta_dict",
+            meta_keys=["image_meta_dict", "label_inverted_meta_dict"],
+            orig_meta_keys="label_meta_dict",
             nearest_interp=True,
-            postfix="inverted",
             to_tensor=[True, False],
             device="cpu",
             num_workers=0 if sys.platform == "darwin" or torch.cuda.is_available() else 2,
@@ -86,10 +89,9 @@ def test_invert(self):
         for d in loader:
             d = inverter(d)
             # this unit test only covers basic function, test_handler_transform_inverter covers more
-            self.assertTupleEqual(d["image"].shape[1:], (1, 100, 100, 100))
             self.assertTupleEqual(d["label"].shape[1:], (1, 100, 100, 100))
             # check the nearest inerpolation mode
-            for i in d["image_inverted"]:
+            for i in d["image"]:
                 torch.testing.assert_allclose(i.to(torch.uint8).to(torch.float), i.to(torch.float))
                 self.assertTupleEqual(i.shape, (1, 100, 101, 107))
             for i in d["label_inverted"]: