step 1 based on PR #4922, adding a lazy transform

Signed-off-by: Wenqi Li <[email protected]>

Author: wyli

docs/source/transforms.rst

@@ -63,6 +63,11 @@ Generic Interfaces
 .. autoclass:: OneOf
     :members:
 
+`LazyTransform`
+^^^^^^^^^^^^^^^
+.. autoclass:: LazyTransform
+    :members:
+
 Vanilla Transforms
 ------------------
 

monai/data/meta_obj.py

@@ -82,6 +82,7 @@ class MetaObj:
     def __init__(self):
         self._meta: dict = MetaObj.get_default_meta()
         self._applied_operations: list = MetaObj.get_default_applied_operations()
+        self._pending_operations: list = MetaObj.get_default_applied_operations()
         self._is_batch: bool = False
 
     @staticmethod
@@ -199,6 +200,19 @@ def push_applied_operation(self, t: Any) -> None:
     def pop_applied_operation(self) -> Any:
         return self._applied_operations.pop()
 
+    @property
+    def pending_operations(self) -> list[dict]:
+        """Get the pending operations. Defaults to ``[]``."""
+        if hasattr(self, "_pending_operations"):
+            return self._pending_operations
+        return MetaObj.get_default_applied_operations()
+
+    def push_pending_operation(self, t: Any) -> None:
+        self._pending_operations.append(t)
+
+    def pop_pending_operation(self) -> Any:
+        return self._pending_operations.pop()
+
     @property
     def is_batch(self) -> bool:
         """Return whether object is part of batch or not."""

monai/data/meta_tensor.py

@@ -23,7 +23,7 @@
 from monai.data.meta_obj import MetaObj, get_track_meta
 from monai.data.utils import affine_to_spacing, decollate_batch, list_data_collate, remove_extra_metadata
 from monai.utils import look_up_option
-from monai.utils.enums import MetaKeys, PostFix, SpaceKeys
+from monai.utils.enums import LazyAttr, MetaKeys, PostFix, SpaceKeys
 from monai.utils.type_conversion import convert_data_type, convert_to_tensor
 
 __all__ = ["MetaTensor"]
@@ -445,6 +445,13 @@ def pixdim(self):
             return [affine_to_spacing(a) for a in self.affine]
         return affine_to_spacing(self.affine)
 
+    def peek_pending_shape(self):
+        """Get the currently expected spatial shape as if all the pending operations are executed."""
+        return self.pending_operations[-1][LazyAttr.SHAPE] if self.pending_operations else self.array.shape[1:]
+
+    def peek_pending_affine(self):
+        return self.pending_operations[-1][LazyAttr.AFFINE] if self.pending_operations else self.affine
+
     def new_empty(self, size, dtype=None, device=None, requires_grad=False):
         """
         must be defined for deepcopy to work

monai/transforms/__init__.py

@@ -449,7 +449,15 @@
     ZoomD,
     ZoomDict,
 )
-from .transform import MapTransform, Randomizable, RandomizableTransform, ThreadUnsafe, Transform, apply_transform
+from .transform import (
+    LazyTransform,
+    MapTransform,
+    Randomizable,
+    RandomizableTransform,
+    ThreadUnsafe,
+    Transform,
+    apply_transform,
+)
 from .utility.array import (
     AddChannel,
     AddCoordinateChannels,

monai/transforms/spatial/array.py

@@ -31,7 +31,7 @@
 from monai.transforms.croppad.array import CenterSpatialCrop, ResizeWithPadOrCrop
 from monai.transforms.intensity.array import GaussianSmooth
 from monai.transforms.inverse import InvertibleTransform
-from monai.transforms.transform import Randomizable, RandomizableTransform, Transform
+from monai.transforms.transform import LazyTransform, Randomizable, RandomizableTransform, Transform
 from monai.transforms.utils import (
     convert_pad_mode,
     create_control_grid,
@@ -48,6 +48,7 @@
     GridSampleMode,
     GridSamplePadMode,
     InterpolateMode,
+    LazyAttr,
     NdimageMode,
     NumpyPadMode,
     SplineMode,
@@ -751,7 +752,7 @@ def inverse(self, data: torch.Tensor) -> torch.Tensor:
         return data
 
 
-class Flip(InvertibleTransform):
+class Flip(InvertibleTransform, LazyTransform):
     """
     Reverses the order of elements along the given spatial axis. Preserves shape.
     See `torch.flip` documentation for additional details:
@@ -771,14 +772,13 @@ class Flip(InvertibleTransform):
     def __init__(self, spatial_axis: Optional[Union[Sequence[int], int]] = None) -> None:
         self.spatial_axis = spatial_axis
 
-    def update_meta(self, img, shape, axes):
+    def update_meta(self, affine, shape, axes):
         # shape and axes include the channel dim
-        affine = img.affine
         mat = convert_to_dst_type(torch.eye(len(affine)), affine)[0]
         for axis in axes:
             sp = axis - 1
             mat[sp, sp], mat[sp, -1] = mat[sp, sp] * -1, shape[axis] - 1
-        img.affine = affine @ mat
+        return affine @ mat
 
     def forward_image(self, img, axes) -> torch.Tensor:
         return torch.flip(img, axes)
@@ -790,9 +790,16 @@ def __call__(self, img: torch.Tensor) -> torch.Tensor:
         """
         img = convert_to_tensor(img, track_meta=get_track_meta())
         axes = map_spatial_axes(img.ndim, self.spatial_axis)
+        if self.lazy_evaluation and isinstance(img, MetaTensor):
+            spatial_chn_shape = [1, *convert_to_numpy(img.peek_pending_shape()).tolist()]
+            affine = img.peek_pending_affine()
+            lazy_affine = self.update_meta(affine, spatial_chn_shape, axes)
+            img.push_pending_operation({LazyAttr.SHAPE: img.peek_pending_shape(), LazyAttr.AFFINE: lazy_affine})
+            self.push_transform(img)
+            return img
         out = self.forward_image(img, axes)
         if get_track_meta():
-            self.update_meta(out, out.shape, axes)
+            out.affine = self.update_meta(out.affine, out.shape, axes)  # type: ignore
             self.push_transform(out)
         return out
 

monai/transforms/transform.py

@@ -26,7 +26,15 @@
 from monai.utils.enums import TransformBackends
 from monai.utils.misc import MONAIEnvVars
 
-__all__ = ["ThreadUnsafe", "apply_transform", "Randomizable", "RandomizableTransform", "Transform", "MapTransform"]
+__all__ = [
+    "ThreadUnsafe",
+    "apply_transform",
+    "LazyTransform",
+    "Randomizable",
+    "RandomizableTransform",
+    "Transform",
+    "MapTransform",
+]
 
 ReturnType = TypeVar("ReturnType")
 
@@ -131,6 +139,26 @@ class ThreadUnsafe:
     pass
 
 
+class LazyTransform:
+    """
+    An interface to denote whether a transform can be applied lazily. It is designed as part of lazy resampling of
+    multiple transforms. Classes inheriting this interface should be able to operate in two modes:
+
+        - ``set_lazy_eval(False)`` (eagerly evaluating), the transform should output the finalized transform
+          results without any pending operations. Both primary data and metadata of the outputs should be up-to-date.
+        - ``set_lazy_eval(True)`` (lazily evaluating), the transform should only execute necessary/lightweight/lossless
+          metadata updates to track any pending operations. The goal is that, in a later stage, the pending operations
+          can be grouped together and evaluated more efficiently and accurately -- each transforms when evaluated
+          independently may cause some information losses.
+
+    """
+
+    lazy_evaluation: bool = False
+
+    def set_lazy_eval(self, value: bool):
+        self.lazy_evaluation = value
+
+
 class Randomizable(ThreadUnsafe):
     """
     An interface for handling random state locally, currently based on a class

monai/utils/__init__.py

@@ -34,6 +34,7 @@
     InterpolateMode,
     InverseKeys,
     JITMetadataKeys,
+    LazyAttr,
     LossReduction,
     MetaKeys,
     Method,

monai/utils/enums.py

@@ -54,6 +54,7 @@
     "AlgoEnsembleKeys",
     "HoVerNetMode",
     "HoVerNetBranch",
+    "LazyAttr",
 ]
 
 
@@ -616,3 +617,16 @@ class HoVerNetBranch(StrEnum):
     HV = "horizontal_vertical"
     NP = "nucleus_prediction"
     NC = "type_prediction"
+
+
+class LazyAttr(StrEnum):
+    """
+    MetaTensor with pending operations requires some key attributes tracked especially when the primary array
+    is not up-to-date due to lazy evaluation.
+    This class specifies the set of key attributes to be tracked for each MetaTensor.
+    """
+
+    SHAPE = "lazy_shape"  # spatial shape
+    AFFINE = "lazy_affine"
+    PADDING_MODE = "lazy_padding_mode"
+    INTERP_MODE = "lazy_interpolation_mode"