Repository: Project-MONAI/MONAI

Issue:
Add `post_func` to `TransformInverter` handler
**Is your feature request related to a problem? Please describe.**
Need to add support to do additional processing after inverting 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
@@ -45,6 +45,7 @@ def __init__(
         collate_fn: Optional[Callable] = no_collation,
         postfix: str = "inverted",
         nearest_interp: Union[bool, Sequence[bool]] = True,
+        post_func: Union[Callable, Sequence[Callable]] = lambda x: x,
         num_workers: Optional[int] = 0,
     ) -> None:
         """
@@ -66,6 +67,8 @@ def __init__(
             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.
+            post_func: post processing for the inverted data, should be a callable function.
+                it also can be a list of callable, each matches to the `output_keys` data.
             num_workers: number of workers when run data loader for inverse transforms,
                 default to 0 as only run one iteration and multi-processing may be even slower.
                 Set to `None`, to use the `num_workers` of the input transform data loader.
@@ -83,6 +86,7 @@ def __init__(
         self.meta_key_postfix = meta_key_postfix
         self.postfix = postfix
         self.nearest_interp = ensure_tuple_rep(nearest_interp, len(self.output_keys))
+        self.post_func = ensure_tuple_rep(post_func, len(self.output_keys))
         self._totensor = ToTensor()
 
     def attach(self, engine: Engine) -> None:
@@ -97,7 +101,9 @@ def __call__(self, engine: Engine) -> None:
         Args:
             engine: Ignite Engine, it can be a trainer, validator or evaluator.
         """
-        for output_key, batch_key, nearest_interp in zip(self.output_keys, self.batch_keys, self.nearest_interp):
+        for output_key, batch_key, nearest_interp, post_funct in zip(
+            self.output_keys, self.batch_keys, self.nearest_interp, self.post_func
+        ):
             transform_key = batch_key + InverseKeys.KEY_SUFFIX
             if transform_key not in engine.state.batch:
                 warnings.warn(f"all the transforms on `{batch_key}` are not InvertibleTransform.")
@@ -121,4 +127,6 @@ def __call__(self, engine: Engine) -> None:
 
             with allow_missing_keys_mode(self.transform):  # type: ignore
                 inverted_key = f"{output_key}_{self.postfix}"
-                engine.state.output[inverted_key] = [self._totensor(i[batch_key]) for i in self.inverter(segs_dict)]
+                engine.state.output[inverted_key] = [
+                    post_funct(self._totensor(i[batch_key])) for i in self.inverter(segs_dict)
+                ]


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
@@ -100,6 +100,7 @@ def _train_func(engine, batch):
             output_keys=["image", "label"],
             batch_keys="image",
             nearest_interp=[True, False],
+            post_func=[lambda x: x + 10, lambda x: x],
             postfix="inverted2",
             num_workers=0 if sys.platform == "darwin" or torch.cuda.is_available() else 2,
         ).attach(engine)