rearrange modules, adds back matrix/grid types

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

Author: wyli

monai/transforms/__init__.py

@@ -228,7 +228,7 @@
 from .io.array import SUPPORTED_READERS, LoadImage, SaveImage
 from .io.dictionary import LoadImaged, LoadImageD, LoadImageDict, SaveImaged, SaveImageD, SaveImageDict
 from .lazy.functional import apply
-from .meta_matrix import matmul
+from .lazy.utils import matmul, resample
 from .meta_utility.dictionary import (
     FromMetaTensord,
     FromMetaTensorD,

monai/transforms/lazy/functional.py

@@ -11,44 +11,15 @@
 
 from typing import Optional, Union
 
-import numpy as np
 import torch
 
 from monai.data.meta_tensor import MetaTensor
 from monai.data.utils import to_affine_nd
-from monai.transforms.meta_matrix import matmul
-from monai.transforms.utility.functional import resample
-from monai.utils import LazyAttr
+from monai.transforms.lazy.utils import is_compatible_kwargs, kwargs_from_pending, mat_from_pending, matmul, resample
 
 __all__ = ["apply"]
 
 
-def mat_from_pending(pending_item):
-    if isinstance(pending_item, (torch.Tensor, np.ndarray)):
-        return pending_item
-    if isinstance(pending_item, dict):
-        return pending_item[LazyAttr.AFFINE]
-    return pending_item
-
-
-def kwargs_from_pending(pending_item):
-    if not isinstance(pending_item, dict):
-        return {}
-    ret = {
-        LazyAttr.INTERP_MODE: pending_item.get(LazyAttr.INTERP_MODE, None),  # interpolation mode
-        LazyAttr.PADDING_MODE: pending_item.get(LazyAttr.PADDING_MODE, None),  # padding mode
-    }
-    if LazyAttr.SHAPE in pending_item:
-        ret[LazyAttr.SHAPE] = pending_item[LazyAttr.SHAPE]
-    if LazyAttr.DTYPE in pending_item:
-        ret[LazyAttr.DTYPE] = pending_item[LazyAttr.DTYPE]
-    return ret
-
-
-def is_compatible_kwargs(kwargs_1, kwargs_2):
-    return True
-
-
 def apply(data: Union[torch.Tensor, MetaTensor], pending: Optional[list] = None):
     """
     This method applies pending transforms to tensors.

monai/transforms/lazy/utils.py

@@ -0,0 +1,119 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import numpy as np
+import torch
+
+import monai
+from monai.config import NdarrayOrTensor
+from monai.utils import LazyAttr
+
+__all__ = ["resample", "matmul"]
+
+
+class Affine:
+    """A class to represent an affine transform matrix."""
+
+    __slots__ = ("data",)
+
+    def __init__(self, data):
+        self.data = data
+
+    @staticmethod
+    def is_affine_shaped(data):
+        """Check if the data is an affine matrix."""
+        if isinstance(data, Affine):
+            return True
+        if isinstance(data, DDF):
+            return False
+        if not hasattr(data, "shape") or len(data.shape) < 2:
+            return False
+        return data.shape[-1] in (3, 4) and data.shape[-2] in (3, 4) and data.shape[-1] == data.shape[-2]
+
+
+class DDF:
+    """A class to represent a dense displacement field."""
+
+    __slots__ = ("data",)
+
+    def __init__(self, data):
+        self.data = data
+
+    @staticmethod
+    def is_ddf_shaped(data):
+        """Check if the data is a DDF."""
+        if isinstance(data, DDF):
+            return True
+        if isinstance(data, Affine):
+            return False
+        if not hasattr(data, "shape") or len(data.shape) < 3:
+            return False
+        return not Affine.is_affine_shaped(data)
+
+
+def matmul(left: torch.Tensor, right: torch.Tensor):
+    if Affine.is_affine_shaped(left) and Affine.is_affine_shaped(right):  # linear transforms
+        if isinstance(left, Affine):
+            left = left.data
+        if isinstance(right, Affine):
+            right = right.data
+        return torch.matmul(left, right)
+    if DDF.is_ddf_shaped(left) and DDF.is_ddf_shaped(right):  # adds DDFs
+        return left + right
+    raise NotImplementedError
+
+
+def mat_from_pending(pending_item):
+    if isinstance(pending_item, (torch.Tensor, np.ndarray)):
+        return pending_item
+    if isinstance(pending_item, dict):
+        return pending_item[LazyAttr.AFFINE]
+    return pending_item
+
+
+def kwargs_from_pending(pending_item):
+    if not isinstance(pending_item, dict):
+        return {}
+    ret = {
+        LazyAttr.INTERP_MODE: pending_item.get(LazyAttr.INTERP_MODE, None),  # interpolation mode
+        LazyAttr.PADDING_MODE: pending_item.get(LazyAttr.PADDING_MODE, None),  # padding mode
+    }
+    if LazyAttr.SHAPE in pending_item:
+        ret[LazyAttr.SHAPE] = pending_item[LazyAttr.SHAPE]
+    if LazyAttr.DTYPE in pending_item:
+        ret[LazyAttr.DTYPE] = pending_item[LazyAttr.DTYPE]
+    return ret
+
+
+def is_compatible_kwargs(kwargs_1, kwargs_2):
+    return True
+
+
+def resample(data: torch.Tensor, matrix: NdarrayOrTensor, kwargs: Optional[dict] = None):
+    """
+    This is a minimal implementation of resample that always uses Affine.
+    """
+    if not Affine.is_affine_shaped(matrix):
+        raise NotImplementedError("calling dense grid resample API not implemented")
+    kwargs = {} if kwargs is None else kwargs
+    init_kwargs = {
+        "spatial_size": kwargs.pop(LazyAttr.SHAPE, data.shape)[1:],
+        "dtype": kwargs.pop(LazyAttr.DTYPE, data.dtype),
+    }
+    call_kwargs = {
+        "mode": kwargs.pop(LazyAttr.INTERP_MODE, None),
+        "padding_mode": kwargs.pop(LazyAttr.PADDING_MODE, None),
+    }
+    resampler = monai.transforms.Affine(affine=matrix, image_only=True, **init_kwargs)
+    with resampler.trace_transform(False):  # don't track this transform in `data`
+        return resampler(img=data, **call_kwargs)

monai/transforms/meta_matrix.py

@@ -890,53 +890,52 @@ def _create_translate(
 
 
 def _create_rotate_90(
-    spatial_dims: int, axis: Tuple[int, int], steps: Optional[int] = 1, eye_func: Callable = np.eye
+    spatial_dims: int, axes: Tuple[int, int], steps: int = 1, eye_func: Callable = np.eye
 ) -> NdarrayOrTensor:
 
     values = [(1, 0, 0, 1), (0, -1, 1, 0), (-1, 0, 0, -1), (0, 1, -1, 0)]
 
     if spatial_dims == 2:
-        if axis != (0, 1):
-            raise ValueError(f"if 'spatial_dims' is 2, 'axis' must be (0, 1) but is {axis}")
+        if axes != (0, 1):
+            raise ValueError(f"if 'spatial_dims' is 2, 'axis' must be (0, 1) but is {axes}")
     elif spatial_dims == 3:
-        if axis not in ((0, 1), (0, 2), (1, 2)):
-            raise ValueError("if 'spatial_dims' is 3, 'axis' must be (0,1), (0, 2), or (1, 2) " f"but is {axis}")
+        if axes not in ((0, 1), (0, 2), (1, 2)):
+            raise ValueError("if 'spatial_dims' is 3, 'axis' must be (0,1), (0, 2), or (1, 2) " f"but is {axes}")
     else:
         raise ValueError(f"'spatial_dims' must be 2 or 3 but is {spatial_dims}")
 
     affine = eye_func(spatial_dims + 1)
 
-    a, b = (0, 1) if spatial_dims == 2 else axis
+    a, b = (0, 1) if spatial_dims == 2 else axes
     affine[a, a], affine[a, b], affine[b, a], affine[b, b] = values[steps % 4]
-    return affine
+    return affine  # type: ignore
 
 
 def create_rotate_90(
     spatial_dims: int,
-    axis: int,
-    steps: Optional[int] = 1,
+    axes: Tuple[int, int] = (0, 1),
+    steps: int = 1,
     device: Optional[torch.device] = None,
     backend: str = TransformBackends.NUMPY,
 ) -> NdarrayOrTensor:
     """
-    create a 2D or 3D rotation matrix
+    create a 2D or 3D rotation90 matrix.
+
     Args:
         spatial_dims: {``2``, ``3``} spatial rank
-        radians: rotation radians
-            when spatial_dims == 3, the `radians` sequence corresponds to
-            rotation in the 1st, 2nd, and 3rd dim respectively.
+        axes: 2 int numbers, defines the plane to rotate with 2 spatial axes.
+            Default: (0, 1), this is the first two axis in spatial dimensions.
+            If axis is negative it counts from the last to the first axis.
+        steps: number of times to rotate by 90 degrees
         device: device to compute and store the output (when the backend is "torch").
         backend: APIs to use, ``numpy`` or ``torch``.
-    Raises:
-        ValueError: When ``radians`` is empty.
-        ValueError: When ``spatial_dims`` is not one of [2, 3].
     """
     _backend = look_up_option(backend, TransformBackends)
     if _backend == TransformBackends.NUMPY:
-        return _create_rotate_90(spatial_dims=spatial_dims, axis=axis, steps=steps, eye_func=np.eye)
+        return _create_rotate_90(spatial_dims=spatial_dims, axes=axes, steps=steps, eye_func=np.eye)
     if _backend == TransformBackends.TORCH:
         return _create_rotate_90(
-            spatial_dims=spatial_dims, axis=axis, steps=steps, eye_func=lambda rank: torch.eye(rank, device=device)
+            spatial_dims=spatial_dims, axes=axes, steps=steps, eye_func=lambda rank: torch.eye(rank, device=device)
         )
     raise ValueError(f"backend {backend} is not supported")
 
@@ -974,9 +973,9 @@ def create_flip(
 ) -> NdarrayOrTensor:
     _backend = look_up_option(backend, TransformBackends)
     if _backend == TransformBackends.NUMPY:
-        return _create_flip(spatial_dims=spatial_dims, spatial_axis=spatial_axis, eye_func=np.eye)
+        return _create_flip(spatial_dims=spatial_dims, spatial_axis=spatial_axis, eye_func=np.eye)  # type: ignore
     if _backend == TransformBackends.TORCH:
-        return _create_flip(
+        return _create_flip(  # type: ignore
             spatial_dims=spatial_dims, spatial_axis=spatial_axis, eye_func=lambda rank: torch.eye(rank, device=device)
         )
     raise ValueError(f"backend {backend} is not supported")

monai/transforms/utility/functional.py

@@ -13,7 +13,7 @@
 
 import torch
 
-from monai.transforms.utility.functional import resample
+from monai.transforms.lazy.functional import resample
 from monai.utils import convert_to_tensor
 from tests.utils import get_arange_img