refactoring

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

Author: wyli

monai/transforms/croppad/array.py

@@ -141,6 +141,26 @@ def _pt_pad(img: torch.Tensor, pad_width, mode, **kwargs) -> torch.Tensor:
         # torch.pad expects `[B, C, H, W, [D]]` shape
         return pad_pt(img.unsqueeze(0), pt_pad_width, mode=mode, **kwargs).squeeze(0)
 
+    @staticmethod
+    def pad_nd(img_t, to_pad_, mode, **kwargs):
+        if mode in {"linear_ramp", "maximum", "mean", "median", "minimum", "symmetric", "empty"}:
+            return Pad._np_pad(img_t, pad_width=to_pad_, mode=mode, **kwargs)
+        mode = convert_pad_mode(dst=img_t, mode=mode).value
+        try:
+            _pad = (
+                Pad._pt_pad
+                if mode in {"reflect", "replicate"}
+                and img_t.dtype not in {torch.int16, torch.int64, torch.bool, torch.uint8}
+                else Pad._np_pad
+            )
+            return _pad(img_t, pad_width=to_pad_, mode=mode, **kwargs)
+        except (ValueError, TypeError, RuntimeError) as err:
+            if isinstance(err, NotImplementedError) or any(
+                k in str(err) for k in ("supported", "unexpected keyword", "implemented")
+            ):
+                return Pad._np_pad(img_t, pad_width=to_pad_, mode=mode, **kwargs)
+            raise ValueError(f"{img_t.shape} {to_pad_} {mode} {kwargs} {img_t.dtype} {img_t.device}") from err
+
     def __call__(  # type: ignore
         self, img: torch.Tensor, to_pad: list[tuple[int, int]] | None = None, mode: str | None = None, **kwargs
     ) -> torch.Tensor:

monai/transforms/croppad/functional.py

@@ -22,7 +22,7 @@
 from monai.data.meta_obj import get_track_meta
 from monai.data.meta_tensor import MetaTensor
 from monai.transforms.inverse import TraceableTransform
-from monai.transforms.utils import convert_pad_mode, create_translate
+from monai.transforms.utils import create_translate
 from monai.utils import TraceKeys, convert_to_dst_type, ensure_tuple
 
 __all__ = ["pad_func", "crop_func"]
@@ -37,7 +37,9 @@ def pad_func(img_t, to_pad_, mode, kwargs, transform_info):
         else torch.eye(4, device=torch.device("cpu"), dtype=torch.float64)
     )
     spatial_rank = max(len(_affine) - 1, 1)
-    if np.asarray(to_pad_).any():
+    if not np.asarray(to_pad_).any():
+        out = img_t
+    else:
         to_pad_ = list(to_pad_)
         if len(to_pad_) < len(img_t.shape):
             to_pad_ = list(to_pad_) + [(0, 0)] * (len(img_t.shape) - len(to_pad_))
@@ -55,27 +57,7 @@ def pad_func(img_t, to_pad_, mode, kwargs, transform_info):
                 extra_info=extra_info,
                 transform_info=transform_info,
             )
-        if mode in {"linear_ramp", "maximum", "mean", "median", "minimum", "symmetric", "empty"}:
-            out = monai.transforms.Pad._np_pad(img_t, pad_width=to_pad_, mode=mode, **kwargs)
-        else:
-            mode = convert_pad_mode(dst=img_t, mode=mode).value
-            try:
-                _pad = (
-                    monai.transforms.Pad._pt_pad
-                    if mode in {"reflect", "replicate"}
-                    and img_t.dtype not in {torch.int16, torch.int64, torch.bool, torch.uint8}
-                    else monai.transforms.Pad._np_pad
-                )
-                out = _pad(img_t, pad_width=to_pad_, mode=mode, **kwargs)
-            except (ValueError, TypeError, RuntimeError) as err:
-                if isinstance(err, NotImplementedError) or any(
-                    k in str(err) for k in ("supported", "unexpected keyword", "implemented")
-                ):
-                    out = monai.transforms.Pad._np_pad(img_t, pad_width=to_pad_, mode=mode, **kwargs)
-                else:
-                    raise ValueError(f"{img_t.shape} {to_pad_} {mode} {kwargs} {img_t.dtype} {img_t.device}") from err
-    else:
-        out = img_t
+        out = monai.transforms.Pad.pad_nd(img_t, to_pad_, mode, **kwargs)
     if get_track_meta():
         to_shift = [-s[0] for s in to_pad_[1:]]  # skipping the channel pad
         out.affine @= convert_to_dst_type(create_translate(spatial_rank, to_shift), _affine)[0]  # type: ignore

monai/transforms/spatial/array.py

@@ -470,7 +470,7 @@ def __call__(
         affine_: np.ndarray
         if affine is not None:
             warnings.warn("arg `affine` is deprecated, the affine of MetaTensor in data_array has higher priority.")
-        input_affine = data_array.affine if isinstance(data_array, MetaTensor) else affine
+        input_affine = data_array.peek_pending_affine() if isinstance(data_array, MetaTensor) else affine
         if input_affine is None:
             warnings.warn("`data_array` is not of type MetaTensor, assuming affine to be identity.")
             # default to identity
@@ -517,7 +517,10 @@ def __call__(
             dtype=dtype,
         )
         if self.recompute_affine and isinstance(data_array, MetaTensor):
-            data_array.affine = scale_affine(affine_, original_spatial_shape, actual_shape)
+            if not self.lazy_evaluation:
+                data_array.affine = scale_affine(affine_, original_spatial_shape, actual_shape)
+            else:
+                raise NotImplementedError("recompute_affine is not supported with lazy evaluation.")
         return data_array
 
     def inverse(self, data: torch.Tensor) -> torch.Tensor: