diff --git a/monai/data/__init__.py b/monai/data/__init__.py index 60574cb565..462ed1c8fa 100644 --- a/monai/data/__init__.py +++ b/monai/data/__init__.py @@ -47,7 +47,7 @@ resolve_writer, ) from .iterable_dataset import CSVIterableDataset, IterableDataset, ShuffleBuffer -from .meta_obj import MetaObj, get_track_meta, get_track_transforms, set_track_meta, set_track_transforms +from .meta_obj import MetaObj, get_track_meta, set_track_meta from .meta_tensor import MetaTensor from .nifti_saver import NiftiSaver from .nifti_writer import write_nifti diff --git a/monai/data/meta_obj.py b/monai/data/meta_obj.py index 3e35a6dd4b..6c56dfe34b 100644 --- a/monai/data/meta_obj.py +++ b/monai/data/meta_obj.py @@ -15,9 +15,8 @@ from typing import Any, Callable, Sequence _TRACK_META = True -_TRACK_TRANSFORMS = True -__all__ = ["get_track_meta", "get_track_transforms", "set_track_meta", "set_track_transforms", "MetaObj"] +__all__ = ["get_track_meta", "set_track_meta", "MetaObj"] def set_track_meta(val: bool) -> None: @@ -26,9 +25,8 @@ def set_track_meta(val: bool) -> None: its data by using subclasses of `MetaObj`. If `False`, then data will be returned with empty metadata. - If both `set_track_meta` and `set_track_transforms` are set to - `False`, then standard data objects will be returned (e.g., `torch.Tensor` and - `np.ndarray`) as opposed to our enhanced objects. + If `set_track_meta` is `False`, then standard data objects will be returned (e.g., + `torch.Tensor` and `np.ndarray`) as opposed to our enhanced objects. By default, this is `True`, and most users will want to leave it this way. However, if you are experiencing any problems regarding metadata, and aren't interested in @@ -38,33 +36,14 @@ def set_track_meta(val: bool) -> None: _TRACK_META = val -def set_track_transforms(val: bool) -> None: - """ - Boolean to set whether transforms are tracked. If `True`, applied transforms will be - associated its data by using subclasses of `MetaObj`. If `False`, then transforms - won't be tracked. - - If both `set_track_meta` and `set_track_transforms` are set to - `False`, then standard data objects will be returned (e.g., `torch.Tensor` and - `np.ndarray`) as opposed to our enhanced objects. - - By default, this is `True`, and most users will want to leave it this way. However, - if you are experiencing any problems regarding transforms, and aren't interested in - preserving transforms, then you can disable it. - """ - global _TRACK_TRANSFORMS - _TRACK_TRANSFORMS = val - - def get_track_meta() -> bool: """ Return the boolean as to whether metadata is tracked. If `True`, metadata will be associated its data by using subclasses of `MetaObj`. If `False`, then data will be returned with empty metadata. - If both `set_track_meta` and `set_track_transforms` are set to - `False`, then standard data objects will be returned (e.g., `torch.Tensor` and - `np.ndarray`) as opposed to our enhanced objects. + If `set_track_meta` is `False`, then standard data objects will be returned (e.g., + `torch.Tensor` and `np.ndarray`) as opposed to our enhanced objects. By default, this is `True`, and most users will want to leave it this way. However, if you are experiencing any problems regarding metadata, and aren't interested in @@ -73,23 +52,6 @@ def get_track_meta() -> bool: return _TRACK_META -def get_track_transforms() -> bool: - """ - Return the boolean as to whether transforms are tracked. If `True`, applied - transforms will be associated its data by using subclasses of `MetaObj`. If `False`, - then transforms won't be tracked. - - If both `set_track_meta` and `set_track_transforms` are set to - `False`, then standard data objects will be returned (e.g., `torch.Tensor` and - `np.ndarray`) as opposed to our enhanced objects. - - By default, this is `True`, and most users will want to leave it this way. However, - if you are experiencing any problems regarding transforms, and aren't interested in - preserving transforms, then you can disable it. - """ - return _TRACK_TRANSFORMS - - class MetaObj: """ Abstract base class that stores data as well as any extra metadata. @@ -177,6 +139,7 @@ def _copy_meta(self, input_objs: list[MetaObj]) -> None: id_in = id(input_objs[0]) if len(input_objs) > 0 else None deep_copy = id(self) != id_in self._copy_attr("meta", input_objs, self.get_default_meta, deep_copy) + self._copy_attr("applied_operations", input_objs, self.get_default_applied_operations, deep_copy) self.is_batch = input_objs[0].is_batch if len(input_objs) > 0 else False def get_default_meta(self) -> dict: @@ -187,6 +150,14 @@ def get_default_meta(self) -> dict: """ return {} + def get_default_applied_operations(self) -> list: + """Get the default applied operations. + + Returns: + default applied operations. + """ + return [] + def __repr__(self) -> str: """String representation of class.""" out: str = super().__repr__() @@ -196,6 +167,14 @@ def __repr__(self) -> str: out += "".join(f"\t{k}: {v}\n" for k, v in self.meta.items()) else: out += "None" + + out += "\nApplied operations\n" + if self.applied_operations is not None: + for i in self.applied_operations: + out += f"\t{str(i)}\n" + else: + out += "None" + out += f"\nIs batch?: {self.is_batch}" return out @@ -210,6 +189,22 @@ def meta(self, d: dict) -> None: """Set the meta.""" self._meta = d + @property + def applied_operations(self) -> list: + """Get the applied operations.""" + return self._applied_operations + + @applied_operations.setter + def applied_operations(self, t: list) -> None: + """Set the applied operations.""" + self._applied_operations = t + + def push_applied_operation(self, t: Any) -> None: + self._applied_operations.append(t) + + def pop_applied_operation(self) -> Any: + return self._applied_operations.pop() + @property def is_batch(self) -> bool: """Return whether object is part of batch or not.""" diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py index d44b780a5e..b63209ed1d 100644 --- a/monai/data/meta_tensor.py +++ b/monai/data/meta_tensor.py @@ -17,7 +17,8 @@ import torch -from monai.data.meta_obj import MetaObj, get_track_meta, get_track_transforms +from monai.config.type_definitions import NdarrayTensor +from monai.data.meta_obj import MetaObj, get_track_meta from monai.data.utils import decollate_batch, list_data_collate from monai.utils.enums import PostFix @@ -72,10 +73,20 @@ class MetaTensor(MetaObj, torch.Tensor): """ @staticmethod - def __new__(cls, x, affine: torch.Tensor | None = None, meta: dict | None = None, *args, **kwargs) -> MetaTensor: + def __new__( + cls, + x, + affine: torch.Tensor | None = None, + meta: dict | None = None, + applied_operations: list | None = None, + *args, + **kwargs, + ) -> MetaTensor: return torch.as_tensor(x, *args, **kwargs).as_subclass(cls) # type: ignore - def __init__(self, x, affine: torch.Tensor | None = None, meta: dict | None = None) -> None: + def __init__( + self, x, affine: torch.Tensor | None = None, meta: dict | None = None, applied_operations: list | None = None + ) -> None: """ If `meta` is given, use it. Else, if `meta` exists in the input tensor, use it. Else, use the default value. Similar for the affine, except this could come from @@ -94,15 +105,24 @@ def __init__(self, x, affine: torch.Tensor | None = None, meta: dict | None = No warnings.warn("Setting affine, but the applied meta contains an affine. This will be overwritten.") self.affine = affine elif "affine" in self.meta: - pass # nothing to do + # by using the setter function, we ensure it is converted to torch.Tensor if not already + self.affine = self.meta["affine"] elif isinstance(x, MetaTensor): self.affine = x.affine else: self.affine = self.get_default_affine() + # applied_operations + if applied_operations is not None: + self.applied_operations = applied_operations + elif isinstance(x, MetaTensor): + self.applied_operations = x.applied_operations + else: + self.applied_operations = self.get_default_applied_operations() # if we are creating a new MetaTensor, then deep copy attributes if isinstance(x, torch.Tensor) and not isinstance(x, MetaTensor): self.meta = deepcopy(self.meta) + self.applied_operations = deepcopy(self.applied_operations) self.affine = self.affine.to(self.device) def _copy_attr(self, attribute: str, input_objs: list[MetaObj], default_fn: Callable, deep_copy: bool) -> None: @@ -126,7 +146,7 @@ def update_meta(rets: Sequence, func, args, kwargs): if not isinstance(ret, MetaTensor): pass # if not tracking, convert to `torch.Tensor`. - elif not (get_track_meta() or get_track_transforms()): + elif not get_track_meta(): ret = ret.as_tensor() # else, handle the `MetaTensor` metadata. else: @@ -221,7 +241,11 @@ def as_dict(self, key: str) -> dict: A dictionary consisting of two keys, the main data (stored under `key`) and the metadata. """ - return {key: self.as_tensor(), PostFix.meta(key): self.meta} + return { + key: self.as_tensor(), + PostFix.meta(key): deepcopy(self.meta), + PostFix.transforms(key): deepcopy(self.applied_operations), + } @property def affine(self) -> torch.Tensor: @@ -229,9 +253,9 @@ def affine(self) -> torch.Tensor: return self.meta["affine"] # type: ignore @affine.setter - def affine(self, d: torch.Tensor) -> None: + def affine(self, d: NdarrayTensor) -> None: """Set the affine.""" - self.meta["affine"] = d + self.meta["affine"] = torch.as_tensor(d, device=self.device) def new_empty(self, size, dtype=None, device=None, requires_grad=False): """ diff --git a/monai/transforms/inverse.py b/monai/transforms/inverse.py index d2e5b2c6ba..bdaa2f9b40 100644 --- a/monai/transforms/inverse.py +++ b/monai/transforms/inverse.py @@ -13,6 +13,7 @@ import torch +from monai.data.meta_tensor import MetaTensor from monai.transforms.transform import Transform from monai.utils.enums import TraceKeys @@ -54,7 +55,8 @@ def trace_key(key: Hashable = None): def push_transform( self, data: Mapping, key: Hashable = None, extra_info: Optional[dict] = None, orig_size: Optional[Tuple] = None ) -> None: - """PUsh to a stack of applied transforms for that key.""" + """Push to a stack of applied transforms for that key.""" + if not self.tracing: return info = {TraceKeys.CLASS_NAME: self.__class__.__name__, TraceKeys.ID: id(self)} @@ -67,17 +69,23 @@ def push_transform( # If class is randomizable transform, store whether the transform was actually performed (based on `prob`) if hasattr(self, "_do_transform"): # RandomizableTransform info[TraceKeys.DO_TRANSFORM] = self._do_transform # type: ignore - # If this is the first, create list - if self.trace_key(key) not in data: - if not isinstance(data, dict): - data = dict(data) - data[self.trace_key(key)] = [] - data[self.trace_key(key)].append(info) + + if key in data and isinstance(data[key], MetaTensor): + data[key].push_applied_operation(info) + else: + # If this is the first, create list + if self.trace_key(key) not in data: + if not isinstance(data, dict): + data = dict(data) + data[self.trace_key(key)] = [] + data[self.trace_key(key)].append(info) def pop_transform(self, data: Mapping, key: Hashable = None): """Remove the most recent applied transform.""" if not self.tracing: return + if key in data and isinstance(data[key], MetaTensor): + return data[key].pop_applied_operation() return data.get(self.trace_key(key), []).pop() @@ -133,7 +141,10 @@ def get_most_recent_transform(self, data: Mapping, key: Hashable = None): """Get most recent transform.""" if not self.tracing: raise RuntimeError("Transform Tracing must be enabled to get the most recent transform.") - transform = data[self.trace_key(key)][-1] + if isinstance(data[key], MetaTensor): + transform = data[key].applied_operations[-1] + else: + transform = data[self.trace_key(key)][-1] self.check_transforms_match(transform) return transform diff --git a/monai/transforms/meta_utility/dictionary.py b/monai/transforms/meta_utility/dictionary.py index 1a9cf4c631..1dcdf3483c 100644 --- a/monai/transforms/meta_utility/dictionary.py +++ b/monai/transforms/meta_utility/dictionary.py @@ -39,7 +39,7 @@ class FromMetaTensord(MapTransform, InvertibleTransform): Dictionary-based transform to convert MetaTensor to a dictionary. If input is `{"a": MetaTensor, "b": MetaTensor}`, then output will - have the form `{"a": torch.Tensor, "a_meta_dict": dict, "b": ...}`. + have the form `{"a": torch.Tensor, "a_meta_dict": dict, "a_transforms": list, "b": ...}`. """ backend = [TransformBackends.TORCH, TransformBackends.NUMPY] @@ -47,9 +47,9 @@ class FromMetaTensord(MapTransform, InvertibleTransform): def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: d = dict(data) for key in self.key_iterator(d): - self.push_transform(d, key) im: MetaTensor = d[key] # type: ignore d.update(im.as_dict(key)) + self.push_transform(d, key) return d def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: @@ -58,8 +58,10 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd # check transform _ = self.get_most_recent_transform(d, key) # do the inverse - im, meta = d[key], d.pop(PostFix.meta(key), None) - im = MetaTensor(im, meta=meta) # type: ignore + im = d[key] + meta = d.pop(PostFix.meta(key), None) + transforms = d.pop(PostFix.transforms(key), None) + im = MetaTensor(im, meta=meta, applied_operations=transforms) # type: ignore d[key] = im # Remove the applied transform self.pop_transform(d, key) @@ -80,8 +82,10 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N d = dict(data) for key in self.key_iterator(d): self.push_transform(d, key) - im, meta = d[key], d.pop(PostFix.meta(key), None) - im = MetaTensor(im, meta=meta) # type: ignore + im = d[key] + meta = d.pop(PostFix.meta(key), None) + transforms = d.pop(PostFix.transforms(key), None) + im = MetaTensor(im, meta=meta, applied_operations=transforms) # type: ignore d[key] = im return d diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py index fb6d922218..87226a585d 100644 --- a/tests/test_meta_tensor.py +++ b/tests/test_meta_tensor.py @@ -22,9 +22,10 @@ from parameterized import parameterized from monai.data import DataLoader, Dataset -from monai.data.meta_obj import get_track_meta, get_track_transforms, set_track_meta, set_track_transforms +from monai.data.meta_obj import get_track_meta, set_track_meta from monai.data.meta_tensor import MetaTensor from monai.data.utils import decollate_batch, list_data_collate +from monai.transforms import BorderPadd, Compose, DivisiblePadd, FromMetaTensord, ToMetaTensord from monai.utils.enums import PostFix from monai.utils.module import pytorch_after from tests.utils import TEST_DEVICES, SkipIfBeforePyTorchVersion, assert_allclose, skip_if_no_cuda @@ -216,10 +217,6 @@ def test_get_set_meta_fns(self): self.assertEqual(get_track_meta(), False) set_track_meta(True) self.assertEqual(get_track_meta(), True) - set_track_transforms(False) - self.assertEqual(get_track_transforms(), False) - set_track_transforms(True) - self.assertEqual(get_track_transforms(), True) @parameterized.expand(TEST_DEVICES) def test_torchscript(self, device): @@ -409,6 +406,71 @@ def test_decollate(self, dtype): self.assertIsInstance(elem, MetaTensor) self.check(elem, im, ids=False) + def test_transforms(self): + key = "im" + _, im = self.get_im() + tr = Compose([BorderPadd(key, 1), DivisiblePadd(key, 16), ToMetaTensord(key), FromMetaTensord(key)]) + num_tr = len(tr.transforms) + data = {key: im, PostFix.meta(key): {"affine": torch.eye(4)}} + + # apply one at a time + is_meta = isinstance(im, MetaTensor) + for i, _tr in enumerate(tr.transforms): + data = _tr(data) + is_meta = isinstance(_tr, ToMetaTensord) + if is_meta: + self.assertEqual(len(data), 1) # im + self.assertIsInstance(data[key], MetaTensor) + n_applied = len(data[key].applied_operations) + else: + self.assertEqual(len(data), 3) # im, im_meta_dict, im_transforms + self.assertIsInstance(data[key], torch.Tensor) + self.assertNotIsInstance(data[key], MetaTensor) + n_applied = len(data[PostFix.transforms(key)]) + + self.assertEqual(n_applied, i + 1) + + # inverse one at a time + is_meta = isinstance(im, MetaTensor) + for i, _tr in enumerate(tr.transforms[::-1]): + data = _tr.inverse(data) + is_meta = isinstance(_tr, FromMetaTensord) + if is_meta: + self.assertEqual(len(data), 1) # im + self.assertIsInstance(data[key], MetaTensor) + n_applied = len(data[key].applied_operations) + else: + self.assertEqual(len(data), 3) # im, im_meta_dict, im_transforms + self.assertIsInstance(data[key], torch.Tensor) + self.assertNotIsInstance(data[key], MetaTensor) + n_applied = len(data[PostFix.transforms(key)]) + + self.assertEqual(n_applied, num_tr - i - 1) + + # apply all in one go + data = tr({key: im, PostFix.meta(key): {"affine": torch.eye(4)}}) + self.assertEqual(len(data), 3) # im, im_meta_dict, im_transforms + self.assertIsInstance(data[key], torch.Tensor) + self.assertNotIsInstance(data[key], MetaTensor) + n_applied = len(data[PostFix.transforms(key)]) + self.assertEqual(n_applied, num_tr) + + # inverse all in one go + data = tr.inverse(data) + self.assertEqual(len(data), 3) # im, im_meta_dict, im_transforms + self.assertIsInstance(data[key], torch.Tensor) + self.assertNotIsInstance(data[key], MetaTensor) + n_applied = len(data[PostFix.transforms(key)]) + self.assertEqual(n_applied, 0) + + def test_construct_with_pre_applied_transforms(self): + key = "im" + _, im = self.get_im() + tr = Compose([BorderPadd(key, 1), DivisiblePadd(key, 16)]) + data = tr({key: im}) + m = MetaTensor(im, applied_operations=data[PostFix.transforms(key)]) + self.assertEqual(len(m.applied_operations), len(tr.transforms)) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_to_from_meta_tensord.py b/tests/test_to_from_meta_tensord.py index 9bbf4592ab..6f46055d6a 100644 --- a/tests/test_to_from_meta_tensord.py +++ b/tests/test_to_from_meta_tensord.py @@ -126,14 +126,12 @@ def test_from_to_meta_tensord(self, device, dtype): self.check(d_dict["m1"], m1.as_tensor(), ids=False) meta_out = {k: v for k, v in d_dict["m1_meta_dict"].items() if k != "affine"} aff_out = d_dict["m1_meta_dict"]["affine"] - self.check(aff_out, m1_aff, ids=True) + self.check(aff_out, m1_aff, ids=False) self.assertEqual(meta_out, m1_meta) # FROM -> inverse d_meta_dict_meta = t_from_meta.inverse(d_dict) - self.assertEqual( - sorted(d_meta_dict_meta.keys()), ["m1", PostFix.transforms("m1"), "m2", PostFix.transforms("m2"), "m3"] - ) + self.assertEqual(sorted(d_meta_dict_meta.keys()), ["m1", "m2", "m3"]) self.check(d_meta_dict_meta["m3"], m3, ids=False) # unchanged (except deep copy in inverse) self.check(d_meta_dict_meta["m1"], m1, ids=False) meta_out = {k: v for k, v in d_meta_dict_meta["m1"].meta.items() if k != "affine"} @@ -143,12 +141,8 @@ def test_from_to_meta_tensord(self, device, dtype): # TO -> Forward t_to_meta = ToMetaTensord(["m1", "m2"]) - del d_dict["m1_transforms"] - del d_dict["m2_transforms"] d_dict_meta = t_to_meta(d_dict) - self.assertEqual( - sorted(d_dict_meta.keys()), ["m1", PostFix.transforms("m1"), "m2", PostFix.transforms("m2"), "m3"] - ) + self.assertEqual(sorted(d_dict_meta.keys()), ["m1", "m2", "m3"]) self.check(d_dict_meta["m3"], m3, ids=True) # unchanged (except deep copy in inverse) self.check(d_dict_meta["m1"], m1, ids=False) meta_out = {k: v for k, v in d_dict_meta["m1"].meta.items() if k != "affine"}