support keep_size=True in lazy Zoom (Project-MONAI#6240)

adds the missing implementation here:

https://github.com/Project-MONAI/MONAI/blob/795bf61adb4c25a03a355fa9c4552473eb763939/monai/transforms/spatial/functional.py#L447-L449

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [x] New tests added to cover the changes.
- [x] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

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

Author: wyli

monai/transforms/spatial/functional.py

@@ -35,6 +35,7 @@
 from monai.transforms.utils import create_rotate, create_translate, scale_affine
 from monai.transforms.utils_pytorch_numpy_unification import allclose
 from monai.utils import (
+    LazyAttr,
     TraceKeys,
     convert_to_dst_type,
     convert_to_numpy,
@@ -432,10 +433,7 @@ def zoom(img, scale_factor, keep_size, mode, padding_mode, align_corners, dtype,
 
     """
     im_shape = img.peek_pending_shape() if isinstance(img, MetaTensor) else img.shape[1:]
-    output_size = [
-        int(math.floor(float(i) * z))
-        for i, z in zip(img.peek_pending_shape() if isinstance(img, MetaTensor) else img.shape[1:], scale_factor)
-    ]
+    output_size = [int(math.floor(float(i) * z)) for i, z in zip(im_shape, scale_factor)]
     xform = scale_affine(im_shape, output_size)
     extra_info = {
         "mode": mode,
@@ -445,9 +443,18 @@ def zoom(img, scale_factor, keep_size, mode, padding_mode, align_corners, dtype,
         "padcrop": {},
     }
     if keep_size:
-        if transform_info.get(TraceKeys.LAZY_EVALUATION, False):
-            raise NotImplementedError("keep_size=True is not supported for lazy evaluation.")
-        output_size = [int(i) for i in img.shape[1:]]
+        do_pad_crop = not np.allclose(output_size, im_shape)
+        if do_pad_crop and transform_info.get(TraceKeys.LAZY_EVALUATION, False):  # update for lazy evaluation
+            _pad_crop = ResizeWithPadOrCrop(spatial_size=im_shape, mode=padding_mode)
+            _pad_crop.lazy_evaluation = True
+            _tmp_img = MetaTensor([], affine=torch.eye(len(output_size) + 1))
+            _tmp_img.push_pending_operation({LazyAttr.SHAPE: list(output_size), LazyAttr.AFFINE: xform})
+            lazy_cropped = _pad_crop(_tmp_img)
+            if isinstance(lazy_cropped, MetaTensor):
+                xform = lazy_cropped.peek_pending_affine()
+                extra_info["padcrop"] = lazy_cropped.pending_operations[-1]
+            extra_info["do_padcrop"] = do_pad_crop
+        output_size = [int(i) for i in im_shape]
     meta_info = TraceableTransform.track_transform_meta(
         img,
         sp_size=output_size,

tests/test_integration_lazy_samples.py

@@ -59,6 +59,9 @@ def run_training_test(root_dir, device="cuda:0", cachedataset=0, readers=(None,
                 keys=["img", "seg"], label_key="seg", spatial_size=[76, 82, 80], pos=1, neg=1, num_samples=4
             ),
             mt.RandRotate90d(keys=["img", "seg"], prob=0.8, spatial_axes=(0, 2)),
+            mt.RandZoomd(
+                keys=["img", "seg"], prob=1.0, min_zoom=1.0, max_zoom=1.0, mode=("trilinear", 0), keep_size=True
+            ),
             mt.ResizeWithPadOrCropD(keys=["img", "seg"], spatial_size=[80, 72, 80]),
             mt.Rotated(keys=["img", "seg"], angle=[np.pi / 2, np.pi / 2, 0], mode="nearest", keep_size=False),
         ],

tests/test_zoom.py

@@ -29,16 +29,25 @@
     test_local_inversion,
 )
 
-VALID_CASES = [(1.5, "nearest", True), (1.5, "nearest", False), (0.8, "bilinear"), (0.8, "area")]
+VALID_CASES = [
+    (1.5, "nearest", True),
+    (1.5, "nearest", False),
+    (0.8, "bilinear"),
+    (0.8, "area"),
+    (1.5, "nearest", False, True),
+    (0.8, "area", False, True),
+]
 
 INVALID_CASES = [((None, None), "bilinear", TypeError), ((0.9, 0.9), "s", ValueError)]
 
 
 class TestZoom(NumpyImageTestCase2D):
     @parameterized.expand(VALID_CASES)
-    def test_pending_ops(self, zoom, mode, align_corners=False):
+    def test_pending_ops(self, zoom, mode, align_corners=False, keep_size=False):
         im = MetaTensor(self.imt[0], meta={"a": "b", "affine": DEFAULT_TEST_AFFINE})
-        zoom_fn = Zoom(zoom=zoom, mode="bilinear", keep_size=False, dtype=torch.float64, align_corners=align_corners)
+        zoom_fn = Zoom(
+            zoom=zoom, mode="bilinear", keep_size=keep_size, dtype=torch.float64, align_corners=align_corners
+        )
         # non-lazy
         expected = zoom_fn(im)
         self.assertIsInstance(expected, MetaTensor)