Invert some keys

monai/transforms/compose.py

@@ -13,6 +13,7 @@
 """
 
 import warnings
+from copy import deepcopy
 from typing import Any, Callable, Optional, Sequence, Union
 
 import numpy as np
@@ -28,6 +29,7 @@
     apply_transform,
 )
 from monai.utils import MAX_SEED, ensure_tuple, get_seed
+from monai.utils.enums import InverseKeys
 
 __all__ = ["Compose"]
 
@@ -156,10 +158,43 @@ def __call__(self, input_):
         return input_
 
     def inverse(self, data):
+        return self.inverse_with_omissions(data, [])
+
+    def inverse_with_omissions(self, data, *to_skip: Sequence[str]):
+        """Perform the inverse on a subset of the applied transformations.
+        This is done by passing a list of strings of the transforms to be skipped.
+
+        We'll get a subset of the transforms, excluding any in the `to_skip` list. We'll
+        also loop across the data and find the applied transforms (e.g., `image_transforms`),
+        and move those into the list of skipped inverse transforms (image_skipped_inverses).
+
+        Args:
+            data: data to be inverted
+            *to_skip: Sequence of class names to skip.
+        """
+        # get all invertible transforms
         invertible_transforms = [t for t in self.flatten().transforms if isinstance(t, InvertibleTransform)]
         if len(invertible_transforms) == 0:
             warnings.warn("inverse has been called but no invertible transforms have been supplied")
 
+        if len(to_skip) > 0:
+            # omit certain invertible transforms
+            invertible_transforms = [t for t in invertible_transforms if t.__class__.__name__ not in to_skip]
+            if len(invertible_transforms) == 0:
+                warnings.warn("all invertible transforms have been omitted")
+
+            data = deepcopy(data)
+            for key in list(data.keys()):
+                inv_key = key + InverseKeys.KEY_SUFFIX
+                if inv_key in data:
+                    to_inverse = [t for t in data[inv_key] if t[InverseKeys.CLASS_NAME] not in to_skip]
+                    to_not_inverse = [t for t in data[inv_key] if t not in to_inverse]
+                    data[key + InverseKeys.KEY_SUFFIX] = to_inverse
+                    skipped_key = key + InverseKeys.KEY_SUFFIX_SKIPPED
+                    if skipped_key not in data:
+                        data[skipped_key] = []
+                    data[key + InverseKeys.KEY_SUFFIX_SKIPPED] += to_not_inverse
+
         # loop backwards over transforms
         for t in reversed(invertible_transforms):
             data = apply_transform(t.inverse, data, self.map_items)

monai/utils/enums.py

@@ -243,6 +243,7 @@ class InverseKeys:
     EXTRA_INFO = "extra_info"
     DO_TRANSFORM = "do_transforms"
     KEY_SUFFIX = "_transforms"
+    KEY_SUFFIX_SKIPPED = "_skipped_inverses"
 
 
 class CommonKeys:

tests/test_inverse_subset.py

@@ -0,0 +1,105 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from typing import TYPE_CHECKING, List, Sequence, Tuple
+from unittest.case import skipUnless
+
+from parameterized import parameterized
+
+from monai.data import create_test_image_2d
+from monai.transforms import AddChanneld, Compose, LoadImaged, RandAxisFlipd, RandFlipd
+from monai.utils import optional_import, set_determinism
+from monai.utils.enums import InverseKeys
+from tests.utils import make_nifti_image
+
+if TYPE_CHECKING:
+    has_nib = True
+else:
+    _, has_nib = optional_import("nibabel")
+
+
+KEYS = ["image", "label"]
+
+TESTS: List[Tuple] = []
+
+
+class ErrorRandAxisFlipd(RandAxisFlipd):
+    def inverse(self, _):
+        raise RuntimeError
+
+
+# remove the ErrorRandAxisFlipd transform. Since its inverse
+# raises an exception, we'll know if this wasn't successful
+TESTS.append(
+    (
+        Compose(
+            [
+                LoadImaged(KEYS),
+                AddChanneld(KEYS),
+                RandAxisFlipd("image"),
+                ErrorRandAxisFlipd("label"),
+                ErrorRandAxisFlipd(KEYS),
+                RandFlipd(KEYS),
+            ]
+        ),
+        "ErrorRandAxisFlipd",
+        (1, 2),
+        False,
+    )
+)
+
+# Nothing is removed, so exception is expected
+TESTS.append(
+    (
+        Compose([LoadImaged(KEYS), AddChanneld(KEYS), ErrorRandAxisFlipd("label")]),
+        "",
+        (0, 0),
+        True,
+    )
+)
+
+
+class TestInverseSubset(unittest.TestCase):
+    def setUp(self):
+        set_determinism(seed=0)
+
+        im_fnames = [make_nifti_image(i) for i in create_test_image_2d(101, 100)]
+        self.data = {k: v for k, v in zip(KEYS, im_fnames)}
+
+    def tearDown(self):
+        set_determinism(seed=None)
+
+    @parameterized.expand(TESTS)
+    @skipUnless(has_nib, "Requires NiBabel")
+    def test_inverse_subset(
+        self,
+        transforms: Compose,
+        to_skip: Sequence[str],
+        expected_num_skipped: Sequence[int],
+        expected_exception: bool,
+    ) -> None:
+        d = transforms(self.data)
+        if not expected_exception:
+            d = transforms.inverse_with_omissions(d, to_skip)
+        else:
+            with self.assertRaises(RuntimeError):
+                d = transforms.inverse_with_omissions(d, to_skip)
+
+        for key, num_skipped in zip(KEYS, expected_num_skipped):
+            if num_skipped > 0:
+                self.assertEqual(len(d[key + InverseKeys.KEY_SUFFIX_SKIPPED]), num_skipped)
+            else:
+                self.assertNotIn(key + InverseKeys.KEY_SUFFIX_SKIPPED, d)
+
+
+if __name__ == "__main__":
+    unittest.main()