Repository: Project-MONAI/MONAI

Issue:
randzoom to partly keep the aspect ratio
**Is your feature request related to a problem? Please describe.**
currently random zoom supports either different random zoom factors for each dimension or the same factor for all dimensions.
would be great to support the use same random factor for some of the dimensions


Discussion:

Resolving patch:
diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py
--- a/monai/transforms/spatial/array.py
+++ b/monai/transforms/spatial/array.py
@@ -751,10 +751,12 @@ class RandZoom(Randomizable, Transform):
             If a float, select a random factor from `[min_zoom, max_zoom]` then apply to all spatial dims
             to keep the original spatial shape ratio.
             If a sequence, min_zoom should contain one value for each spatial axis.
+            If 2 values provided for 3D data, use the first value for both H & W dims to keep the same zoom ratio.
         max_zoom: Max zoom factor. Can be float or sequence same size as image.
             If a float, select a random factor from `[min_zoom, max_zoom]` then apply to all spatial dims
             to keep the original spatial shape ratio.
             If a sequence, max_zoom should contain one value for each spatial axis.
+            If 2 values provided for 3D data, use the first value for both H & W dims to keep the same zoom ratio.
         mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``}
             The interpolation mode. Defaults to ``"area"``.
             See also: https://pytorch.org/docs/stable/nn.functional.html#interpolate
@@ -788,14 +790,11 @@ def __init__(
         self.keep_size = keep_size
 
         self._do_transform = False
-        self._zoom: Union[List[float], float] = 1.0
+        self._zoom: Sequence[float] = [1.0]
 
     def randomize(self, data: Optional[Any] = None) -> None:
         self._do_transform = self.R.random_sample() < self.prob
         self._zoom = [self.R.uniform(l, h) for l, h in zip(self.min_zoom, self.max_zoom)]
-        if len(self._zoom) == 1:
-            # to keep the spatial shape ratio, use same random zoom factor for all dims
-            self._zoom = self._zoom[0]
 
     def __call__(
         self,
@@ -823,6 +822,12 @@ def __call__(
         _dtype = np.float32
         if not self._do_transform:
             return img.astype(_dtype)
+        if len(self._zoom) == 1:
+            # to keep the spatial shape ratio, use same random zoom factor for all dims
+            self._zoom = ensure_tuple_rep(self._zoom[0], img.ndim - 1)
+        elif len(self._zoom) == 2 and img.ndim > 3:
+            # if 2 zoom factors provided for 3D data, use the first factor for H and W dims, second factor for D dim
+            self._zoom = ensure_tuple_rep(self._zoom[0], img.ndim - 2) + ensure_tuple(self._zoom[-1])
         zoomer = Zoom(self._zoom, keep_size=self.keep_size)
         return zoomer(
             img,
diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py
--- a/monai/transforms/spatial/dictionary.py
+++ b/monai/transforms/spatial/dictionary.py
@@ -932,10 +932,12 @@ class RandZoomd(Randomizable, MapTransform):
             If a float, select a random factor from `[min_zoom, max_zoom]` then apply to all spatial dims
             to keep the original spatial shape ratio.
             If a sequence, min_zoom should contain one value for each spatial axis.
+            If 2 values provided for 3D data, use the first value for both H & W dims to keep the same zoom ratio.
         max_zoom: Max zoom factor. Can be float or sequence same size as image.
             If a float, select a random factor from `[min_zoom, max_zoom]` then apply to all spatial dims
             to keep the original spatial shape ratio.
             If a sequence, max_zoom should contain one value for each spatial axis.
+            If 2 values provided for 3D data, use the first value for both H & W dims to keep the same zoom ratio.
         mode: {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``}
             The interpolation mode. Defaults to ``"area"``.
             See also: https://pytorch.org/docs/stable/nn.functional.html#interpolate
@@ -974,14 +976,11 @@ def __init__(
         self.keep_size = keep_size
 
         self._do_transform = False
-        self._zoom: Union[Sequence[float], float]
+        self._zoom: Sequence[float] = [1.0]
 
     def randomize(self, data: Optional[Any] = None) -> None:
         self._do_transform = self.R.random_sample() < self.prob
         self._zoom = [self.R.uniform(l, h) for l, h in zip(self.min_zoom, self.max_zoom)]
-        if len(self._zoom) == 1:
-            # to keep the spatial shape ratio, use same random zoom factor for all dims
-            self._zoom = self._zoom[0]
 
     def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
         # match the spatial dim of first item
@@ -989,6 +988,14 @@ def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.nda
         d = dict(data)
         if not self._do_transform:
             return d
+
+        img_dims = data[self.keys[0]].ndim
+        if len(self._zoom) == 1:
+            # to keep the spatial shape ratio, use same random zoom factor for all dims
+            self._zoom = ensure_tuple_rep(self._zoom[0], img_dims - 1)
+        elif len(self._zoom) == 2 and img_dims > 3:
+            # if 2 zoom factors provided for 3D data, use the first factor for H and W dims, second factor for D dim
+            self._zoom = ensure_tuple_rep(self._zoom[0], img_dims - 2) + ensure_tuple(self._zoom[-1])
         zoomer = Zoom(self._zoom, keep_size=self.keep_size)
         for idx, key in enumerate(self.keys):
             d[key] = zoomer(


Test patch:
diff --git a/tests/test_rand_zoom.py b/tests/test_rand_zoom.py
--- a/tests/test_rand_zoom.py
+++ b/tests/test_rand_zoom.py
@@ -61,6 +61,20 @@ def test_invalid_inputs(self, _, min_zoom, max_zoom, mode, raises):
             random_zoom = RandZoom(prob=1.0, min_zoom=min_zoom, max_zoom=max_zoom, mode=mode)
             random_zoom(self.imt[0])
 
+    def test_auto_expand_3d(self):
+        random_zoom = RandZoom(
+            prob=1.0,
+            min_zoom=[0.8, 0.7],
+            max_zoom=[1.2, 1.3],
+            mode="nearest",
+            keep_size=False,
+        )
+        random_zoom.set_random_state(1234)
+        test_data = np.random.randint(0, 2, size=[2, 2, 3, 4])
+        zoomed = random_zoom(test_data)
+        np.testing.assert_allclose(random_zoom._zoom, (1.048844, 1.048844, 0.962637), atol=1e-2)
+        np.testing.assert_allclose(zoomed.shape, (2, 2, 3, 3))
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_rand_zoomd.py b/tests/test_rand_zoomd.py
--- a/tests/test_rand_zoomd.py
+++ b/tests/test_rand_zoomd.py
@@ -58,6 +58,21 @@ def test_invalid_inputs(self, _, min_zoom, max_zoom, mode, raises):
             random_zoom = RandZoomd(key, prob=1.0, min_zoom=min_zoom, max_zoom=max_zoom, mode=mode)
             random_zoom({key: self.imt[0]})
 
+    def test_auto_expand_3d(self):
+        random_zoom = RandZoomd(
+            keys="img",
+            prob=1.0,
+            min_zoom=[0.8, 0.7],
+            max_zoom=[1.2, 1.3],
+            mode="nearest",
+            keep_size=False,
+        )
+        random_zoom.set_random_state(1234)
+        test_data = {"img": np.random.randint(0, 2, size=[2, 2, 3, 4])}
+        zoomed = random_zoom(test_data)
+        np.testing.assert_allclose(random_zoom._zoom, (1.048844, 1.048844, 0.962637), atol=1e-2)
+        np.testing.assert_allclose(zoomed["img"].shape, (2, 2, 3, 3))
+
 
 if __name__ == "__main__":
     unittest.main()