Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Lazy eval refactor #6257

Closed
wants to merge 141 commits into from
Closed
Changes from 1 commit
Commits
Show all changes
141 commits
Select commit Hold shift + click to select a range
06d5842
Initial commit
atbenmurray Mar 27, 2023
2507a0b
Refactors and simplification of api / calls
atbenmurray Mar 28, 2023
56b618e
Refactoring functional spatial transform to support lazy_evaluation p…
atbenmurray Mar 28, 2023
315e5b7
Making lazy_evaluation parameter consistent across spatial array lazy…
atbenmurray Mar 28, 2023
4e69313
Updating apply and dictionary spatial transforms to support
atbenmurray Mar 29, 2023
99fa35b
Missing changes to spatial dictionary
atbenmurray Mar 29, 2023
a798e0e
Enable call-time overriding of lazy_evaluation in croppad/functional and
atbenmurray Mar 29, 2023
114b8cc
Updating croppad dictionary and croppad array for call-time laziness
atbenmurray Mar 29, 2023
5df4261
Fix for push_transform; updating the lazy integration test
atbenmurray Mar 29, 2023
84d65f7
removing issue.txt
atbenmurray Mar 30, 2023
7ec442f
. Renaming apply_transform to apply_pending as it is more semantically
atbenmurray Mar 31, 2023
8f62776
Tweaked array ApplyPending and ApplyPendingd for array and dict data
atbenmurray Mar 31, 2023
2492223
Added missing license text.
atbenmurray Mar 31, 2023
0064326
Missing import for apply_pending
atbenmurray Mar 31, 2023
0646602
Merge branch 'dev' into lazy_eval_refactor
atbenmurray Apr 14, 2023
05eb3f9
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 14, 2023
788cd74
Fixing syntax error created by merge of dev to PR
atbenmurray Apr 14, 2023
4bd22f8
Fixing issues with pad_func post merge
atbenmurray Apr 14, 2023
1e0510c
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 14, 2023
c36b69f
Fixing issues in crop while running unit tests
atbenmurray Apr 14, 2023
8a4f310
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 14, 2023
3c69776
Fixing issues with OneOf/SomeOf/RandomOrder after dev merge
atbenmurray Apr 14, 2023
6b94aa1
More croppad fixes post dev merge
atbenmurray Apr 14, 2023
bf1c01a
Fixing test_nvtx_decorator after def merge
atbenmurray Apr 14, 2023
19745f9
Removed use of **kwargs parameter in tests as it isn't supported by a…
atbenmurray Apr 23, 2023
b219f13
Setting use of lazy_evaluation to False for Rand2D/Rand3DElastic
atbenmurray Apr 23, 2023
6af8f24
Fixed test_spatial_pad
atbenmurray Apr 23, 2023
173c161
Renaming lazy_evaluation to lazy
atbenmurray Apr 23, 2023
2262ba9
Fixed all outstanding usages of apply_pending in tests. Removed the
atbenmurray Apr 23, 2023
aaaea13
Merge branch 'dev' into lazy_eval_refactor
atbenmurray Apr 24, 2023
088d14e
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 24, 2023
3e6a637
Accepting formatting auto-fixes
atbenmurray Apr 24, 2023
98a53bf
Renaming surviving lazy_evaluation instances to lazy.
atbenmurray Apr 24, 2023
28a0175
Renamed the last errant 'lazy_evaluation' instances to 'lazy'
atbenmurray Apr 24, 2023
4e9b6c6
Fixing ruff errors for over-long line lengths
atbenmurray Apr 24, 2023
ad23bea
Fix for issue when the 'lazy' property is set to None
atbenmurray Apr 24, 2023
81549e2
Fix for _apply_transform when lazy is None, post merge tidy up
atbenmurray Apr 24, 2023
77eff5b
Removing the redundant check for log_stats from
atbenmurray Apr 24, 2023
7c53d26
Further formatting auto-fixes
atbenmurray Apr 24, 2023
ee5378f
Making default for overrides be None instead of {}.
atbenmurray Apr 24, 2023
28c05d0
Changed execute_pending_transforms to not make a shallow copy of dict…
atbenmurray Apr 24, 2023
d96fef0
Fixing docstring for Compose
atbenmurray Apr 24, 2023
654cf48
Fixing type complaints
atbenmurray Apr 24, 2023
64ec242
Fix for tests/tests_integration_lazy_samples. The modified lazy
atbenmurray Apr 25, 2023
f38eb9b
Flake8 fix
atbenmurray Apr 25, 2023
14349ff
Flake8 fixes
atbenmurray Apr 25, 2023
51fda45
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 25, 2023
c5742c0
Autoformat style fixes
atbenmurray Apr 25, 2023
ef5173e
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 25, 2023
916b0a9
Fixing bug introduced by flake8 fix
atbenmurray Apr 25, 2023
20c1cfc
Merge branch 'dev' into lazy_eval_refactor
atbenmurray Apr 25, 2023
c94341f
Removal of obsolete lazy eval code
atbenmurray Apr 25, 2023
cc51f2f
Autoformat shuffle
atbenmurray Apr 25, 2023
06aad01
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 25, 2023
3ae46e9
Resolving mypy issues
atbenmurray Apr 25, 2023
95a6a5f
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 25, 2023
81559fa
Finally got to grips with mypy troubles
atbenmurray Apr 25, 2023
af9cad1
Final mypy complaints
atbenmurray Apr 25, 2023
a9c8849
Chasing autofixes again
atbenmurray Apr 25, 2023
aadbae3
test_invert_warn_pending is now obsolete due to this no longer being an
atbenmurray Apr 25, 2023
1cccf9b
Adding documentation for ApplyPending/ApplyPendingd. Adding more
atbenmurray Apr 25, 2023
9ae1140
More autofix
atbenmurray Apr 25, 2023
bb2ad0f
Restoring lazy.setters for croppad and spatial lazy transforms
atbenmurray Apr 25, 2023
5abe294
Undoing the previous commit until I get to the bottom of the randaffined
atbenmurray Apr 25, 2023
c2bf4a7
Making apis for OneOf/SomeOf/RandomOrder consistent with Compose
atbenmurray Apr 26, 2023
b16a87f
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 26, 2023
816666e
Removing partial docstring for 'lazy', which isn't on the init parameter
atbenmurray Apr 26, 2023
8318121
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 26, 2023
0e10c24
Adding execute_pending_transforms to lazy/functional __all__
atbenmurray Apr 26, 2023
f323ff9
Removing LazyMode; lazy is now True / False / None
atbenmurray Apr 26, 2023
4b69ec6
Renamed execute_pending_transforms to apply_pending_transforms.
atbenmurray Apr 26, 2023
29ce2b7
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 26, 2023
6774aaa
Removed "LazyMode" from enums.py __all__
atbenmurray Apr 26, 2023
de5b9b8
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 26, 2023
1c7d172
Modified ApplyPending/ApplyPendingd in the light of the override issue
atbenmurray Apr 26, 2023
bb22c13
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 26, 2023
2159ce8
Correction to type info for execute_compose lazy parameter.
atbenmurray Apr 26, 2023
2c3172d
Removal of comment in docstring about IdentityD being used to cause
atbenmurray Apr 26, 2023
bb1214c
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 26, 2023
5a7820c
Logging refactored with pending pipeline logging enabled
atbenmurray Apr 26, 2023
503db85
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 26, 2023
1af35fd
All logging_names are now None by default to prevent massive spam in
atbenmurray Apr 26, 2023
75771a8
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 26, 2023
d97095a
Merge branch 'dev' into lazy_eval_refactor
atbenmurray Apr 26, 2023
af6cd62
More autofixes
atbenmurray Apr 26, 2023
cd246a8
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 26, 2023
1403d77
Making apply_transform logger_name None by default
atbenmurray Apr 26, 2023
5c4c10a
Added a dictionary test to the compose logging test suite
atbenmurray Apr 26, 2023
5755a65
Updated docstring for _apply_transform
atbenmurray Apr 27, 2023
ca0905a
Removed spurious empty string at the end of one of the logging test case
atbenmurray Apr 27, 2023
e8fdf5a
All spatial transforms now initialize LazyTransform correctly
atbenmurray Apr 27, 2023
70d3517
Adding setters for lazy back in.
atbenmurray Apr 27, 2023
b361a13
Merge branch 'dev' into lazy_eval_refactor
atbenmurray Apr 27, 2023
a6ad51e
Sorting out line length complaint in compose logging tests
atbenmurray Apr 27, 2023
7ec8429
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 27, 2023
305b88e
Fix for old lazy_evaluation name creeping in during merge
atbenmurray Apr 27, 2023
2aa0ada
Autofix
atbenmurray Apr 27, 2023
7b092ca
Merge branch 'Project-MONAI:dev' into lazy_eval_refactor
atbenmurray Apr 27, 2023
62431ff
Using get_logger so that named loggers all get set up correctly
atbenmurray Apr 27, 2023
6cdf4c6
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray Apr 27, 2023
2ea89c1
Fixes for logging; additional compose logging tests; missing imports
atbenmurray Apr 28, 2023
628a7f8
Adding ApplyPending and ApplyPendingd to transforms.rst
atbenmurray Apr 28, 2023
5b34755
Fixed transforms.rst modifications
atbenmurray Apr 28, 2023
d547b81
Refactor of lazy to provide in_order and out_of_order lazy execution
atbenmurray May 2, 2023
03386dc
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray May 2, 2023
d2fa521
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 2, 2023
eb641c6
Tidied up __all__ issues caused by refactor
atbenmurray May 2, 2023
d7d7cdb
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray May 2, 2023
29b37bf
Removed can_invert flag from compose compiler policy
atbenmurray May 2, 2023
3194980
Moved to use of a TraceKey.STATUS flag. This can be used to track
atbenmurray May 3, 2023
5ff516d
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 3, 2023
eebb3de
Formatting autofixes
atbenmurray May 3, 2023
d966856
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray May 3, 2023
e39e0a3
Lazy {reorder: lazy_last_nosync} now passes
atbenmurray May 3, 2023
cf58b9f
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 3, 2023
883f8c5
Fixed broken call to isinstance in lazy executors
atbenmurray May 4, 2023
e81f359
Autofix
atbenmurray May 4, 2023
db1486b
Resolving lint issue
atbenmurray May 4, 2023
b151607
Missing passing logger_name through to _apply_transform when mapping …
atbenmurray May 4, 2023
24286f1
Autofix
atbenmurray May 4, 2023
43f7bc0
Added missing lazy=None option from integration tests for lazy
atbenmurray May 4, 2023
b217dbe
Autofix
atbenmurray May 4, 2023
ed1d952
Fixing lint issues
atbenmurray May 5, 2023
e7f3a07
Autofixes
atbenmurray May 5, 2023
0751c52
lint fix
atbenmurray May 5, 2023
6114a67
started on the lazy resampling topic page
atbenmurray May 5, 2023
240059b
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray May 5, 2023
23ff89c
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 5, 2023
5e1b81d
Removing unused Sequence import from lazy/executors
atbenmurray May 6, 2023
5dd6113
Renaming ComposeCompiler to ExecuteOptions in test_compose
atbenmurray May 6, 2023
74bed34
Update tests/test_compose.py
atbenmurray May 10, 2023
76e4b85
Execution options is no longer a poltergeist. Refactored is_tensor_in…
atbenmurray May 10, 2023
1c4fa6c
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray May 10, 2023
8dde40d
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 10, 2023
7fd32e2
Removing unnecessary active_key property
atbenmurray May 10, 2023
10d95ce
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray May 10, 2023
0d6a847
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 10, 2023
1ceeb1a
Fix for multi-sampled data for transforms that aceept *args; refactor…
atbenmurray May 11, 2023
38bdef6
Merge branch 'lazy_eval_refactor' of github.com:atbenmurray/monai int…
atbenmurray May 11, 2023
6e0cfd9
Renaming LazyTrait partially_lazy -> checks_data
atbenmurray May 21, 2023
f2fb06a
Merge branch 'Project-MONAI:dev' into lazy_eval_refactor
atbenmurray May 21, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Refactor of lazy to provide in_order and out_of_order lazy execution
Signed-off-by: Ben Murray <ben.murray@gmail.com>
  • Loading branch information
atbenmurray committed May 2, 2023
commit d547b81eed9c588348d58a7bb2ffa178e5e8ddb2
2 changes: 1 addition & 1 deletion docs/source/transforms.rst
Original file line number Diff line number Diff line change
@@ -1927,7 +1927,7 @@ Lazy (Dict)
^^^^^^^^^^^

`ApplyPendingd`
""""""""""""""
"""""""""""""""

.. autoclass:: ApplyPendingd
:members:
229 changes: 212 additions & 17 deletions monai/transforms/compose.py
Original file line number Diff line number Diff line change
@@ -24,10 +24,13 @@
import monai
from monai.apps.utils import get_logger
from monai.config import NdarrayOrTensor
from monai.data.meta_tensor import MetaTensor
from monai.transforms.inverse import InvertibleTransform
from monai.transforms.lazy.array import ApplyPending
from monai.transforms.lazy.dictionary import ApplyPendingd

# For backwards compatibility (so this still works: from monai.transforms.compose import MapTransform)
from monai.transforms.lazy.functional import apply_pending_transforms
from monai.transforms.lazy.executors import apply_pending_transforms
from monai.transforms.traits import LazyTrait, ThreadUnsafe
from monai.transforms.transform import ( # noqa: F401
LazyTransform,
@@ -52,6 +55,7 @@ def execute_compose(
start: int = 0,
end: int | None = None,
lazy: bool | None = False,
lazy_strategy: str = "in_order",
overrides: dict | None = None,
threading: bool = False,
logger_name: str | None = None,
@@ -75,6 +79,9 @@ def execute_compose(
lazy: whether to enable lazy evaluation for lazy transforms. If False, transforms will be
carried out on a transform by transform basis. If True, all lazy transforms will
be executed by accumulating changes and resampling as few times as possible.
lazy_strategy: this field controls how execution occurs when processing data lazily. Permitted
options are "in_order", "out_of_order". Please see `Compose`_ for more details of what these
options mean. In general, you should not need to change this from its default.
overrides: this optional parameter allows you to specify a dictionary of parameters that should be overridden
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to specify here we could accept a dict or a nested dict here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I'll create an example

when executing a pipeline. These each parameter that is compatible with a given transform is then applied
to that transform before it is executed. Note that overrides are currently only applied when lazy
@@ -109,12 +116,138 @@ def execute_compose(
if threading:
_transform = deepcopy(_transform) if isinstance(_transform, ThreadUnsafe) else _transform
data = apply_transform(
_transform, data, map_items, unpack_items, lazy=lazy, overrides=overrides, logger_name=logger_name
_transform,
data,
map_items,
unpack_items,
lazy=lazy,
lazy_strategy=lazy_strategy,
overrides=overrides,
logger_name=logger_name,
)
data = apply_pending_transforms(data, overrides, logger_name=logger_name)
data = apply_pending_transforms(data, None, overrides, logger_name=logger_name)
return data


class ComposeCompiler:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As we've discussed today at the meeting I'd suggest something like LazyPolicy given the current state of this class. If the intent is to evolve this to do more complicated work to "compile" in some manner we might want to put that into a new class anyway and keep something to do policy like this separate.

"""
The ComposeCompiler is an implementation class that is required to parse options for Compose. It should currently
be considered an implementation detail that should not be interacted with directly by users of MONAI, although that
may change in subsequent releases. Its job is to parse options provided to `Compose.__call__`_ to set execution
modes for lazy resampling.

See `Compose`_ for a detailed explanation of lazy resampling.
"""

def __init__(self):
# construct the list of options
options = {"reorder": {"lazy_last": self.reorder_lazy_last, "lazy_last_nosync": self.reorder_lazy_last_nosync}}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would avoid having nested structures like this, instead:

Suggested change
options = {"reorder": {"lazy_last": self.reorder_lazy_last, "lazy_last_nosync": self.reorder_lazy_last_nosync}}
options = {("reorder","lazy_last"): self.reorder_lazy_last, ("reorder","lazy_last_nosync"): self.reorder_lazy_last_nosync}

Also the string keywords here should be in enums and the structure that's expected should be documented in the docstring.

self.options = options

def __call__(self, transforms, lazy: bool | None, options: dict | None = None):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

options is only ever used with one key-value pair so maybe just have two variables here. Otherwise eval_options or something, with more docstring on what it means.

"""
Get a policy object that controls the way `Compose`_ executes a list of transforms.

Args:
transforms: a list of transforms to be executed
lazy: the current lazy mode (False, None, or True)
options: the options that determine the execution policy

Returns:
a dictionary specifying the execution policy

"""
if lazy is False or options is None:
return ComposeCompiler.generate_policy()

if len(options.keys()) > 1:
raise ValueError("Only one option can currently be set")

for k, v in options.items():
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

k, v = monai.utils.first(options.items())?

if k not in self.options.keys():
raise KeyError(
f"'{k}' is not a valid option key. Valid options are " f"{tuple(k for k in self.options.keys())}"
)

option = self.options[k]
if v not in option.keys():
raise KeyError(f"'{v}' is not a valid option value. Value options for " f"'{k}' are {option.keys()}")

action = option[v]

return action(transforms=transforms, lazy=lazy)

@classmethod
def reorder_lazy_last(cls, *, transforms: list, lazy: bool | None, **kwargs):
subsections = list()
subsection_starts = list()
# pass 1: split the transform list into subsections
i_s = 0
for i_t in range(len(transforms)):
if isinstance(transforms[i_t], (ApplyPending, ApplyPendingd)):
# this subsection ends and is added to the subsection list
if i_s < i_t:
subsections.append(transforms[i_s:i_t])
subsection_starts.append(i_s)
# add apply pending in its own list
subsections.append([transforms[i_t]])
subsection_starts.append(i_t)
i_s = i_t + 1

if i_s != len(transforms):
subsections.append(transforms[i_s:])
subsection_starts.append(i_s)

# pass 2: calculate the permuted indices
permuted_indices = list()
for sub_start, subsection in zip(subsection_starts, subsections):
for i_s, s in enumerate(subsection):
if not cls._executing_lazily(s, lazy):
permuted_indices.append(i_s + sub_start)
for i_s, s in enumerate(subsection):
if cls._executing_lazily(s, lazy):
permuted_indices.append(i_s + sub_start)

# pass 2: sort the subsections
reordered = list()
for subsection in subsections:
# non-lazy, lazy
subsection = [t for t in subsection if not cls._executing_lazily(t, lazy)] + [
t for t in subsection if cls._executing_lazily(t, lazy)
]
reordered.extend(subsection)

return ComposeCompiler.generate_policy({"indices": permuted_indices})

@classmethod
def reorder_lazy_last_nosync(cls, *, transforms: list, **_):
"""
'reorder: lazy_last_nosync' is implemented through use of the 'out_of_order' execution
policy. See 'Compose'_ for details of this policy.
Args:
transforms: Not used by this method

Returns:

"""
return cls.generate_policy({"can_invert": False, "lazy_policy": "out_of_order"})

@staticmethod
def generate_policy(overrides: dict | None = None):
default_policy = {"indices": None, "transforms": None, "can_invert": True, "lazy_policy": "in_order"}
if overrides is not None:
for k, v in overrides.items():
default_policy[k] = v
return default_policy

@staticmethod
def _executing_lazily(t, lazy_policy):
if isinstance(t, LazyTrait):
lazy_ = t.lazy if lazy_policy is None else lazy_policy
return lazy_
return False


class Compose(Randomizable, InvertibleTransform):
"""
``Compose`` provides the ability to chain a series of callables together in
@@ -262,6 +395,7 @@ def __init__(
unpack_items: bool = False,
lazy: bool | None = False,
overrides: dict | None = None,
options: dict | None = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please add documentation for this options then I guess we can start reviewing the PR?

logger_name: str | None = None,
) -> None:
if transforms is None:
@@ -272,6 +406,7 @@ def __init__(
self.set_random_state(seed=get_seed())
self.lazy = lazy
self.overrides = overrides
self.options = options
self.logger_name = logger_name

def set_random_state(self, seed: int | None = None, state: np.random.RandomState | None = None) -> Compose:
@@ -350,33 +485,94 @@ def __len__(self):
return len(self.flatten().transforms)

def __call__(self, input_, start=0, end=None, threading=False, lazy: bool | None = None):
return execute_compose(
lazy_ = self.lazy if lazy is None else lazy
policy = ComposeCompiler()(self.transforms, lazy_)
lazy_strategy = policy["lazy_policy"]
indices = policy["indices"]

# permute the transforms if required
transforms = self.transforms if indices is None else [self.transforms[i] for i in indices]

result = execute_compose(
input_,
self.transforms,
transforms,
start=start,
end=end,
map_items=self.map_items,
unpack_items=self.unpack_items,
lazy=self.lazy, # type: ignore
lazy_strategy=lazy_strategy,
overrides=self.overrides,
threading=threading,
logger_name=self.logger_name,
)

# if the transforms were permuted, record it in the metadata for inversion
if indices is not None:
if isinstance(result, monai.data.MetaTensor):
self.push_transform(result, extra_info={"applied_order": indices})
elif isinstance(result, Mapping):
for key in result: # dictionary not change size during iteration
if isinstance(result[key], monai.data.MetaTensor) or self.trace_key(key) in result:
self.push_transform(result, key, extra_info={"applied_order": indices})

return result

def inverse(self, data):
invertible_transforms = [t for t in self.flatten().transforms if isinstance(t, InvertibleTransform)]
if not invertible_transforms:
warnings.warn("inverse has been called but no invertible transforms have been supplied")
if self.options is not None:
compiler = ComposeCompiler()
policy = compiler(self.transforms, self.lazy, self.options)
else:
policy = ComposeCompiler().generate_policy()

indices = policy["indices"]
can_invert = policy["can_invert"]
if can_invert is False:
raise ValueError("'inverse' is not supported with options {self.options}")

data_ = deepcopy(data)

if indices is not None:
applied_order = None
if isinstance(data_, monai.data.MetaTensor):
applied_order = self.pop_transform(data_)[TraceKeys.EXTRA_INFO]["applied_order"]
elif isinstance(data_, Mapping):
for key in data_:
if isinstance(data_[key], monai.data.MetaTensor) or self.trace_key(key) in data_:
applied_order = self.pop_transform(data_, key)[TraceKeys.EXTRA_INFO]["applied_order"]
else:
raise RuntimeError(
f"Inverse only implemented for Mapping (dictionary) or MetaTensor data, got type {type(data)}."
)
if applied_order is None:
# no invertible transforms have been applied
return data_

# loop backwards over transforms
for o in reversed(applied_order):
if isinstance(self.transforms[o], InvertibleTransform):
data_ = apply_transform(self.transforms[o].inverse, data_, self.map_items, self.unpack_items)
else:
invertible_transforms = [t for t in self.flatten().transforms if isinstance(t, InvertibleTransform)]
if not invertible_transforms:
warnings.warn("inverse has been called but no invertible transforms have been supplied")

# loop backwards over transforms
for t in reversed(invertible_transforms):
if isinstance(t, LazyTrait) and t.lazy:
if self.lazy is not False:
warnings.warn(
f"inversing {t.__class__.__name__} lazily may not implemented"
"please set `lazy=False` before calling inverse."
f"'lazy' is set to {self.lazy} but lazy execution is not supported when inverting. "
f"'lazy' has been overridden to False for the call to inverse"
)
data = apply_transform(t.inverse, data, self.map_items, self.unpack_items)
return data
# loop backwards over transforms
for t in reversed(invertible_transforms):
# if isinstance(t, LazyTrait) and t.lazy:
# warnings.warn(
# f"inversing {t.__class__.__name__} lazily may not implemented"
# "please set `lazy=False` before calling inverse."
# )
data_ = apply_transform(
t.inverse, data_, self.map_items, self.unpack_items, lazy=False, logger_name=self.logger_name
)
return data_


class OneOf(Compose):
@@ -759,8 +955,7 @@ def inverse(self, data):

# loop backwards over transforms
for o in reversed(applied_order):
transform = self.transforms[o]
if isinstance(transform, InvertibleTransform):
if isinstance(self.transforms[o], InvertibleTransform):
data = apply_transform(self.transforms[o].inverse, data, self.map_items, self.unpack_items)

return data
4 changes: 2 additions & 2 deletions monai/transforms/inverse.py
Original file line number Diff line number Diff line change
@@ -23,7 +23,7 @@
from monai.data.meta_obj import MetaObj, get_track_meta
from monai.data.meta_tensor import MetaTensor
from monai.data.utils import to_affine_nd
from monai.transforms.traits import LazyTrait
from monai.transforms.traits import InvertibleTrait, LazyTrait
from monai.transforms.transform import Transform
from monai.utils import LazyAttr, MetaKeys, TraceKeys, convert_to_dst_type, convert_to_numpy, convert_to_tensor

@@ -325,7 +325,7 @@ def trace_transform(self, to_trace: bool):
self.tracing = prev


class InvertibleTransform(TraceableTransform):
class InvertibleTransform(TraceableTransform, InvertibleTrait):
"""Classes for invertible transforms.

This class exists so that an ``invert`` method can be implemented. This allows, for
4 changes: 2 additions & 2 deletions monai/transforms/lazy/array.py
Original file line number Diff line number Diff line change
@@ -11,12 +11,12 @@

from __future__ import annotations

from monai.transforms.inverse import InvertibleTransform
from monai.transforms.traits import InvertibleTrait

__all__ = ["ApplyPending"]


class ApplyPending(InvertibleTransform):
class ApplyPending(InvertibleTrait):
"""
ApplyPending can be inserted into a pipeline that is being executed lazily in order to ensure
resampling happens before the next transform. It doesn't do anything itself, but its presence
12 changes: 5 additions & 7 deletions monai/transforms/lazy/dictionary.py
Original file line number Diff line number Diff line change
@@ -13,13 +13,12 @@
from __future__ import annotations

from monai.config import KeysCollection
from monai.transforms.inverse import InvertibleTransform
from monai.transforms.transform import MapTransform
from monai.transforms.traits import InvertibleTrait, MapTrait

__all__ = ["ApplyPendingd", "ApplyPendingD", "ApplyPendingDict"]


class ApplyPendingd(InvertibleTransform, MapTransform):
class ApplyPendingd(InvertibleTrait, MapTrait):
"""
ApplyPendingd can be inserted into a pipeline that is being executed lazily in order
to ensure resampling happens before the next transform. It doesn't do anything itself,
@@ -28,12 +27,11 @@ class ApplyPendingd(InvertibleTransform, MapTransform):
See ``Compose`` for a detailed explanation of the lazy resampling feature.

Args:
keys: the keys on which the transform operates. As of 1.2, this field must be set
but it doesn't alter the behaviour of lazy resampling.
keys: the keys for tensors that should have their pending transforms executed
"""

def __init__(self, keys: KeysCollection):
super().__init__(keys, True)
def __init__(self, keys: KeysCollection | None):
self.keys = keys

def __call__(self, data):
if not isinstance(data, dict):
Loading