backward compatible inverse

monai/apps/detection/transforms/dictionary.py

@@ -136,7 +136,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
         return d
 
     def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
         for key in self.key_iterator(d):
             tr = self.get_most_recent_transform(d, key)
             src_mode, dst_mode = tr[TraceKeys.EXTRA_INFO]["src"], tr[TraceKeys.EXTRA_INFO]["dst"]
@@ -191,7 +191,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
         return d
 
     def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
         for key in self.key_iterator(d):
             tr = self.get_most_recent_transform(d, key)
             original_mode = tr[TraceKeys.EXTRA_INFO]["mode"]
@@ -280,7 +280,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
         return d
 
     def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
         for key in self.key_iterator(d):
             transform = self.get_most_recent_transform(d, key)
             affine = transform["extra_info"]["affine"]
@@ -594,7 +594,7 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torc
         return d
 
     def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
 
         for key in self.key_iterator(d):
             transform = self.get_most_recent_transform(d, key)
@@ -675,7 +675,7 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torc
         return d
 
     def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
 
         for key in self.key_iterator(d):
             transform = self.get_most_recent_transform(d, key)
@@ -752,7 +752,7 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torc
         return d
 
     def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
 
         for key in self.key_iterator(d):
             transform = self.get_most_recent_transform(d, key)
@@ -1274,7 +1274,7 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Mapping[Hashable, t
         return d
 
     def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
 
         for key in self.key_iterator(d):
             transform = self.get_most_recent_transform(d, key)
@@ -1360,7 +1360,7 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Mapping[Hashable, t
         return d
 
     def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torch.Tensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
         if self._rand_k % 4 == 0:
             return d
 

monai/config/__init__.py

@@ -11,6 +11,7 @@
 
 from .deviceconfig import (
     USE_COMPILED,
+    USE_META_DICT,
     IgniteInfo,
     get_config_values,
     get_gpu_info,

monai/config/deviceconfig.py

@@ -27,6 +27,8 @@
 except (OptionalImportError, ImportError, AttributeError):
     HAS_EXT = USE_COMPILED = False
 
+USE_META_DICT = os.environ.get("USE_META_DICT", "0") == "1"  # set to True for compatibility, use meta dict.
+
 psutil, has_psutil = optional_import("psutil")
 psutil_version = psutil.__version__ if has_psutil else "NOT INSTALLED or UNKNOWN VERSION."
 
@@ -38,6 +40,7 @@
     "print_gpu_info",
     "print_debug_info",
     "USE_COMPILED",
+    "USE_META_DICT",
     "IgniteInfo",
 ]
 
@@ -89,7 +92,7 @@ def print_config(file=sys.stdout):
     """
     for k, v in get_config_values().items():
         print(f"{k} version: {v}", file=file, flush=True)
-    print(f"MONAI flags: HAS_EXT = {HAS_EXT}, USE_COMPILED = {USE_COMPILED}")
+    print(f"MONAI flags: HAS_EXT = {HAS_EXT}, USE_COMPILED = {USE_COMPILED}, USE_META_DICT = {USE_META_DICT}")
     print(f"MONAI rev id: {monai.__revision_id__}")
     print(f"MONAI __file__: {monai.__file__}")
 

monai/data/__init__.py

@@ -71,6 +71,7 @@
 from .thread_buffer import ThreadBuffer, ThreadDataLoader
 from .torchscript_utils import load_net_with_metadata, save_net_with_metadata
 from .utils import (
+    PICKLE_KEY_SUFFIX,
     affine_to_spacing,
     compute_importance_map,
     compute_shape_offset,

monai/data/utils.py

@@ -27,6 +27,7 @@
 import torch
 from torch.utils.data._utils.collate import default_collate
 
+from monai import config
 from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor, PathLike
 from monai.data.meta_obj import MetaObj
 from monai.networks.layers.simplelayers import GaussianFilter
@@ -90,6 +91,7 @@
     "remove_keys",
     "remove_extra_metadata",
     "get_extra_metadata_keys",
+    "PICKLE_KEY_SUFFIX",
 ]
 
 # module to be used by `torch.save`
@@ -391,6 +393,32 @@ def dev_collate(batch, level: int = 1, logger_name: str = "dev_collate"):
     return
 
 
+PICKLE_KEY_SUFFIX = TraceKeys.KEY_SUFFIX
+
+
+def pickle_operations(data, key=PICKLE_KEY_SUFFIX, is_encode: bool = True):
+    """
+    Applied_operations are dictionaries with varying sizes, this method converts them to bytes so that we can (de-)collate.
+
+    Args:
+        data: a list or dictionary with substructures to be pickled/unpickled.
+        key: the key suffix for the target substructures, defaults to "_transforms" (`data.utils.PICKLE_KEY_SUFFIX`).
+        is_encode: whether it's encoding using pickle.dumps (True) or decoding using pickle.loads (False).
+    """
+    if isinstance(data, Mapping):
+        data = dict(data)
+        for k in data:
+            if f"{k}".endswith(key):
+                if is_encode and not isinstance(data[k], bytes):
+                    data[k] = pickle.dumps(data[k], 0)
+                if not is_encode and isinstance(data[k], bytes):
+                    data[k] = pickle.loads(data[k])
+        return {k: pickle_operations(v, key=key, is_encode=is_encode) for k, v in data.items()}
+    elif isinstance(data, (list, tuple)):
+        return [pickle_operations(item, key=key, is_encode=is_encode) for item in data]
+    return data
+
+
 def collate_meta_tensor(batch):
     """collate a sequence of meta tensor sequences/dictionaries into
     a single batched metatensor or a dictionary of batched metatensor"""
@@ -423,6 +451,8 @@ def list_data_collate(batch: Sequence):
     data = [i for k in batch for i in k] if isinstance(elem, list) else batch
     key = None
     try:
+        if config.USE_META_DICT:
+            data = pickle_operations(data)  # bc 0.9.0
         if isinstance(elem, Mapping):
             ret = {}
             for k in elem:
@@ -573,10 +603,16 @@ def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None):
             deco[k] = [deepcopy(deco[k]) for _ in range(b)]
     if isinstance(deco, Mapping):
         _gen = zip_longest(*deco.values(), fillvalue=fill_value) if pad else zip(*deco.values())
-        return [dict(zip(deco, item)) for item in _gen]
+        ret = [dict(zip(deco, item)) for item in _gen]
+        if not config.USE_META_DICT:
+            return ret
+        return pickle_operations(ret, is_encode=False)  # bc 0.9.0
     if isinstance(deco, Iterable):
         _gen = zip_longest(*deco, fillvalue=fill_value) if pad else zip(*deco)
-        return [list(item) for item in _gen]
+        ret_list = [list(item) for item in _gen]
+        if not config.USE_META_DICT:
+            return ret_list
+        return pickle_operations(ret_list, is_encode=False)  # bc 0.9.0
     raise NotImplementedError(f"Unable to de-collate: {batch}, type: {type(batch)}.")
 
 

monai/transforms/__init__.py

@@ -257,6 +257,7 @@
     Activations,
     AsDiscrete,
     FillHoles,
+    Invert,
     KeepLargestConnectedComponent,
     LabelFilter,
     LabelToContour,
@@ -574,6 +575,7 @@
 from .utils import (
     Fourier,
     allow_missing_keys_mode,
+    attach_hook,
     compute_divisible_spatial_size,
     convert_applied_interp_mode,
     convert_pad_mode,
@@ -607,6 +609,7 @@
     rescale_array_int_max,
     rescale_instance_array,
     resize_center,
+    sync_meta_info,
     weighted_patch_samples,
     zero_margins,
 )

monai/transforms/inverse.py

@@ -16,6 +16,7 @@
 
 import torch
 
+from monai import transforms
 from monai.data.meta_tensor import MetaTensor
 from monai.transforms.transform import Transform
 from monai.utils.enums import TraceKeys
@@ -250,6 +251,22 @@ class InvertibleTransform(TraceableTransform):
 
     """
 
+    def inverse_update(self, data):
+        """
+        This function is to be called before every `self.inverse(data)`,
+        update each MetaTensor `data[key]` using `data[key_transforms]` and `data[key_meta_dict]`,
+        for MetaTensor backward compatibility 0.9.0.
+        """
+        if not isinstance(data, dict) or not isinstance(self, transforms.MapTransform):
+            return data
+        d = dict(data)
+        for k in self.key_iterator(data):
+            transform_key = transforms.TraceableTransform.trace_key(k)
+            if transform_key not in data or not data[transform_key]:
+                continue
+            d = transforms.sync_meta_info(k, data, t=False)
+        return d
+
     def inverse(self, data: Any) -> Any:
         """
         Inverse of ``__call__``.

monai/transforms/meta_utility/dictionary.py

@@ -15,7 +15,6 @@
 Class names are ended with 'd' to denote dictionary-based transforms.
 """
 
-from copy import deepcopy
 from typing import Dict, Hashable, Mapping
 
 from monai.config.type_definitions import NdarrayOrTensor
@@ -53,7 +52,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
         return d
 
     def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
         for key in self.key_iterator(d):
             # check transform
             _ = self.get_most_recent_transform(d, key)
@@ -90,7 +89,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
         return d
 
     def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
-        d = deepcopy(dict(data))
+        d = dict(data)
         for key in self.key_iterator(d):
             # check transform
             _ = self.get_most_recent_transform(d, key)

monai/transforms/post/array.py

@@ -21,10 +21,17 @@
 
 from monai.config.type_definitions import NdarrayOrTensor
 from monai.data.meta_obj import get_track_meta
+from monai.data.meta_tensor import MetaTensor
 from monai.networks import one_hot
 from monai.networks.layers import GaussianFilter, apply_filter
+from monai.transforms.inverse import InvertibleTransform
 from monai.transforms.transform import Transform
-from monai.transforms.utils import fill_holes, get_largest_connected_component_mask, get_unique_labels
+from monai.transforms.utils import (
+    convert_applied_interp_mode,
+    fill_holes,
+    get_largest_connected_component_mask,
+    get_unique_labels,
+)
 from monai.transforms.utils_pytorch_numpy_unification import unravel_index
 from monai.utils import (
     TransformBackends,
@@ -46,6 +53,7 @@
     "MeanEnsemble",
     "ProbNMS",
     "VoteEnsemble",
+    "Invert",
 ]
 
 
@@ -765,3 +773,45 @@ def __call__(self, prob_map: NdarrayOrTensor):
             prob_map[slices] = 0
 
         return outputs
+
+
+class Invert(Transform):
+    """
+    Utility transform to automatically invert the previously applied transforms.
+    """
+
+    def __init__(
+        self,
+        transform: Optional[InvertibleTransform] = None,
+        nearest_interp: Union[bool, Sequence[bool]] = True,
+        device: Union[Union[str, torch.device], Sequence[Union[str, torch.device]]] = "cpu",
+        post_func: Union[Callable, Sequence[Callable]] = lambda x: x,
+    ) -> None:
+        """
+        Args:
+            transform: the previously applied transform.
+            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.
+            device: move the inverted results to a target device before `post_func`, default to "cpu".
+            post_func: postprocessing for the inverted MetaTensor, should be a callable function.
+        """
+        if not isinstance(transform, InvertibleTransform):
+            raise ValueError("transform is not invertible, can't invert transform for the data.")
+        self.transform = transform
+        self.nearest_interp = nearest_interp
+        self.device = device
+        self.post_func = post_func
+
+    def __call__(self, data):
+        if not isinstance(data, MetaTensor):
+            return data
+
+        if self.nearest_interp:
+            data.applied_operations = convert_applied_interp_mode(
+                trans_info=data.applied_operations, mode="nearest", align_corners=None
+            )
+
+        data = data.detach()
+        inverted = self.transform.inverse(data)
+        inverted = self.post_func(inverted.to(self.device))
+        return inverted

monai/transforms/post/dictionary.py

@@ -21,6 +21,7 @@
 
 import torch
 
+from monai import config
 from monai.config.type_definitions import KeysCollection, NdarrayOrTensor, PathLike
 from monai.data.csv_saver import CSVSaver
 from monai.data.meta_tensor import MetaTensor
@@ -39,8 +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_applied_interp_mode
-from monai.utils import convert_to_tensor, deprecated_arg, ensure_tuple, ensure_tuple_rep
-from monai.utils.enums import PostFix
+from monai.utils import PostFix, convert_to_tensor, deprecated_arg, ensure_tuple, ensure_tuple_rep
 
 __all__ = [
     "ActivationsD",
@@ -634,12 +634,13 @@ def __call__(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]:
                     warnings.warn(f"transform info of `{orig_key}` is not available or no InvertibleTransform applied.")
                     continue
 
+            orig_meta_key = orig_meta_key or f"{orig_key}_{meta_key_postfix}"
             if orig_key in d and isinstance(d[orig_key], MetaTensor):
                 transform_info = d[orig_key].applied_operations
                 meta_info = d[orig_key].meta
             else:
                 transform_info = d[InvertibleTransform.trace_key(orig_key)]
-                meta_info = d.get(orig_meta_key or f"{orig_key}_{meta_key_postfix}", {})
+                meta_info = d.get(orig_meta_key, {})
             if nearest_interp:
                 transform_info = convert_applied_interp_mode(
                     trans_info=deepcopy(transform_info), mode="nearest", align_corners=None
@@ -656,7 +657,9 @@ def __call__(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]:
 
             # construct the input dict data
             input_dict = {orig_key: inputs}
-
+            if config.USE_META_DICT:
+                input_dict[InvertibleTransform.trace_key(orig_key)] = transform_info
+                input_dict[PostFix.meta(orig_key)] = meta_info
             with allow_missing_keys_mode(self.transform):  # type: ignore
                 inverted = self.transform.inverse(input_dict)
 
@@ -665,8 +668,9 @@ def __call__(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, Any]:
                 inverted_data = self._totensor(inverted[orig_key])
             else:
                 inverted_data = inverted[orig_key]
+                if config.USE_META_DICT and InvertibleTransform.trace_key(orig_key) in d:
+                    d[InvertibleTransform.trace_key(orig_key)] = inverted_data.applied_operations
             d[key] = post_func(inverted_data.to(device))
-
             # save the inverted meta dict
             if orig_meta_key in d:
                 meta_key = meta_key or f"{key}_{meta_key_postfix}"