evaluate cases

Signed-off-by: Wenqi Li <wenqil@nvidia.com>

Author: wyli

monai/transforms/compose.py

@@ -49,7 +49,7 @@ def eval_lazy_stack(
     if not lazy_evaluation:
         return data  # eager evaluation
     if isinstance(data, monai.data.MetaTensor):
-        if not isinstance(upcoming, LazyTransform):
+        if not (isinstance(upcoming, LazyTransform) and upcoming.lazy_evaluation):
             data, _ = mt.apply_transforms(data, mode=mode, padding_mode=padding_mode)
         return data
     if isinstance(data, Mapping):

monai/transforms/croppad/array.py

@@ -800,9 +800,9 @@ def __init__(
         self.padder = Pad(mode=mode, **pad_kwargs)
 
     @Crop.lazy_evaluation.setter  # type: ignore
-    def lazy_evaluation(self, val: bool):
-        self.lazy_evaluation = val
-        self.padder.lazy_evaluation = val
+    def lazy_evaluation(self, _val: bool):
+        self._lazy_evaluation = False  # foreground can't be computed lazily
+        self.padder.lazy_evaluation = False
 
     def compute_bounding_box(self, img: torch.Tensor):
         """
@@ -839,10 +839,8 @@ def crop_pad(
         ret = self.padder.__call__(img=cropped, to_pad=pad_width, mode=mode, **pad_kwargs)
         # combine the traced cropping and padding into one transformation
         # by taking the padded info and placing it in a key inside the crop info.
-        if get_track_meta():
-            ret_: MetaTensor = ret  # type: ignore
-            app_op = ret_.applied_operations.pop(-1)
-            ret_.applied_operations[-1][TraceKeys.EXTRA_INFO]["pad_info"] = app_op
+        if get_track_meta() and isinstance(ret, MetaTensor):
+            ret.applied_operations[-1][TraceKeys.EXTRA_INFO]["pad_info"] = ret.applied_operations.pop()
         return ret
 
     def __call__(self, img: torch.Tensor, mode: str | None = None, **pad_kwargs):  # type: ignore

monai/transforms/croppad/dictionary.py

@@ -684,6 +684,11 @@ def __init__(
         super().__init__(keys, cropper=cropper, allow_missing_keys=allow_missing_keys)
         self.mode = ensure_tuple_rep(mode, len(self.keys))
 
+    @LazyTransform.lazy_evaluation.setter  # type: ignore
+    def lazy_evaluation(self, val: bool):
+        self._lazy_evaluation = False  # foreground can't be computed lazily
+        self.cropper.lazy_evaluation = False
+
     def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> dict[Hashable, torch.Tensor]:
         d = dict(data)
         self.cropper: CropForeground

monai/transforms/inverse.py

@@ -163,11 +163,11 @@ def track_transform_tensor(
             orig_affine = convert_to_dst_type(orig_affine, affine)[0]
             affine = orig_affine @ to_affine_nd(len(orig_affine) - 1, affine, dtype=affine.dtype)
             out_obj.meta[MetaKeys.AFFINE] = convert_to_tensor(affine, device=torch.device("cpu"))
-        if (
-            not isinstance(data_t, MetaTensor)
-            or not get_track_meta()
-            or not transform_info
-            or not transform_info.get(TraceKeys.TRACING)
+        if not (
+            isinstance(data_t, MetaTensor)
+            and get_track_meta()
+            and transform_info
+            and transform_info.get(TraceKeys.TRACING)
         ):
             if key is not None:
                 data[key] = data_t.copy_meta_from(out_obj) if isinstance(data_t, MetaTensor) else data_t

monai/transforms/lazy/functional.py

@@ -70,5 +70,7 @@ def apply_transforms(
     data = resample(data, cumulative_xform, sp_size, cur_kwargs)
     if isinstance(data, MetaTensor):
         for p in pending:
+            for attr in LazyAttr:
+                p.pop(attr, None)
             data.push_applied_operation(p)
     return data, pending