Apply and Resample

.github/workflows/pythonapp.yml

@@ -88,6 +88,10 @@ jobs:
       name: Install torch cpu from pytorch.org (Windows only)
       run: |
         python -m pip install torch==1.12.1+cpu torchvision==0.13.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+    - if: runner.os == 'Linux'
+      name: Install itk pre-release (Linux only)
+      run: |
+        python -m pip install --pre -U itk
     - name: Install the dependencies
       run: |
         python -m pip install torch==1.12.1 torchvision==0.13.1

docs/source/transforms.rst

@@ -22,11 +22,31 @@ Generic Interfaces
     :members:
     :special-members: __call__
 
+`RandomizableTrait`
+^^^^^^^^^^^^^^^^^^^
+.. autoclass:: RandomizableTrait
+    :members:
+
+`LazyTrait`
+^^^^^^^^^^^
+.. autoclass:: LazyTrait
+    :members:
+
+`MultiSampleTrait`
+^^^^^^^^^^^^^^^^^^
+.. autoclass:: MultiSampleTrait
+    :members:
+
 `Randomizable`
 ^^^^^^^^^^^^^^
 .. autoclass:: Randomizable
     :members:
 
+`LazyTransform`
+^^^^^^^^^^^^^^^
+.. autoclass:: LazyTransform
+    :members:
+
 `RandomizableTransform`
 ^^^^^^^^^^^^^^^^^^^^^^^
 .. autoclass:: RandomizableTransform

monai/data/image_reader.py

@@ -318,7 +318,12 @@ def _get_meta_dict(self, img) -> Dict:
 
         """
         img_meta_dict = img.GetMetaDataDictionary()
-        meta_dict = {key: img_meta_dict[key] for key in img_meta_dict.GetKeys() if not key.startswith("ITK_")}
+        meta_dict = {}
+        for key in img_meta_dict.GetKeys():
+            if key.startswith("ITK_"):
+                continue
+            val = img_meta_dict[key]
+            meta_dict[key] = np.asarray(val) if type(val).__name__.startswith("itk") else val
 
         meta_dict["spacing"] = np.asarray(img.GetSpacing())
         return meta_dict

monai/data/meta_obj.py

@@ -82,6 +82,7 @@ class MetaObj:
     def __init__(self):
         self._meta: dict = MetaObj.get_default_meta()
         self._applied_operations: list = MetaObj.get_default_applied_operations()
+        self._pending_operations: list = MetaObj.get_default_applied_operations()  # the same default as applied_ops
         self._is_batch: bool = False
 
     @staticmethod
@@ -199,6 +200,19 @@ def push_applied_operation(self, t: Any) -> None:
     def pop_applied_operation(self) -> Any:
         return self._applied_operations.pop()
 
+    @property
+    def pending_operations(self) -> list[dict]:
+        """Get the pending operations. Defaults to ``[]``."""
+        if hasattr(self, "_pending_operations"):
+            return self._pending_operations
+        return MetaObj.get_default_applied_operations()  # the same default as applied_ops
+
+    def push_pending_operation(self, t: Any) -> None:
+        self._pending_operations.append(t)
+
+    def pop_pending_operation(self) -> Any:
+        return self._pending_operations.pop()
+
     @property
     def is_batch(self) -> bool:
         """Return whether object is part of batch or not."""

monai/data/meta_tensor.py

@@ -23,8 +23,8 @@
 from monai.data.meta_obj import MetaObj, get_track_meta
 from monai.data.utils import affine_to_spacing, decollate_batch, list_data_collate, remove_extra_metadata
 from monai.utils import look_up_option
-from monai.utils.enums import MetaKeys, PostFix, SpaceKeys
-from monai.utils.type_conversion import convert_data_type, convert_to_tensor
+from monai.utils.enums import LazyAttr, MetaKeys, PostFix, SpaceKeys
+from monai.utils.type_conversion import convert_data_type, convert_to_numpy, convert_to_tensor
 
 __all__ = ["MetaTensor"]
 
@@ -445,6 +445,20 @@ def pixdim(self):
             return [affine_to_spacing(a) for a in self.affine]
         return affine_to_spacing(self.affine)
 
+    def peek_pending_shape(self):
+        """Get the currently expected spatial shape as if all the pending operations are executed."""
+        res = None
+        if self.pending_operations:
+            res = self.pending_operations[-1].get(LazyAttr.SHAPE, None)
+        # default to spatial shape (assuming channel-first input)
+        return tuple(convert_to_numpy(self.shape, wrap_sequence=True).tolist()[1:]) if res is None else res
+
+    def peek_pending_affine(self):
+        res = None
+        if self.pending_operations:
+            res = self.pending_operations[-1].get(LazyAttr.AFFINE, None)
+        return self.affine if res is None else res
+
     def new_empty(self, size, dtype=None, device=None, requires_grad=False):
         """
         must be defined for deepcopy to work

monai/networks/nets/attentionunet.py

@@ -143,12 +143,27 @@ def forward(self, g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
 
 
 class AttentionLayer(nn.Module):
-    def __init__(self, spatial_dims: int, in_channels: int, out_channels: int, submodule: nn.Module, dropout=0.0):
+    def __init__(
+        self,
+        spatial_dims: int,
+        in_channels: int,
+        out_channels: int,
+        submodule: nn.Module,
+        up_kernel_size=3,
+        strides=2,
+        dropout=0.0,
+    ):
         super().__init__()
         self.attention = AttentionBlock(
             spatial_dims=spatial_dims, f_g=in_channels, f_l=in_channels, f_int=in_channels // 2
         )
-        self.upconv = UpConv(spatial_dims=spatial_dims, in_channels=out_channels, out_channels=in_channels, strides=2)
+        self.upconv = UpConv(
+            spatial_dims=spatial_dims,
+            in_channels=out_channels,
+            out_channels=in_channels,
+            strides=strides,
+            kernel_size=up_kernel_size,
+        )
         self.merge = Convolution(
             spatial_dims=spatial_dims, in_channels=2 * in_channels, out_channels=in_channels, dropout=dropout
         )
@@ -174,7 +189,7 @@ class AttentionUnet(nn.Module):
         channels (Sequence[int]): sequence of channels. Top block first. The length of `channels` should be no less than 2.
         strides (Sequence[int]): stride to use for convolutions.
         kernel_size: convolution kernel size.
-        upsample_kernel_size: convolution kernel size for transposed convolution layers.
+        up_kernel_size: convolution kernel size for transposed convolution layers.
         dropout: dropout ratio. Defaults to no dropout.
     """
 
@@ -210,9 +225,9 @@ def __init__(
         )
         self.up_kernel_size = up_kernel_size
 
-        def _create_block(channels: Sequence[int], strides: Sequence[int], level: int = 0) -> nn.Module:
+        def _create_block(channels: Sequence[int], strides: Sequence[int]) -> nn.Module:
             if len(channels) > 2:
-                subblock = _create_block(channels[1:], strides[1:], level=level + 1)
+                subblock = _create_block(channels[1:], strides[1:])
                 return AttentionLayer(
                     spatial_dims=spatial_dims,
                     in_channels=channels[0],
@@ -227,17 +242,19 @@ def _create_block(channels: Sequence[int], strides: Sequence[int], level: int =
                         ),
                         subblock,
                     ),
+                    up_kernel_size=self.up_kernel_size,
+                    strides=strides[0],
                     dropout=dropout,
                 )
             else:
                 # the next layer is the bottom so stop recursion,
-                # create the bottom layer as the sublock for this layer
-                return self._get_bottom_layer(channels[0], channels[1], strides[0], level=level + 1)
+                # create the bottom layer as the subblock for this layer
+                return self._get_bottom_layer(channels[0], channels[1], strides[0])
 
         encdec = _create_block(self.channels, self.strides)
         self.model = nn.Sequential(head, encdec, reduce_channels)
 
-    def _get_bottom_layer(self, in_channels: int, out_channels: int, strides: int, level: int) -> nn.Module:
+    def _get_bottom_layer(self, in_channels: int, out_channels: int, strides: int) -> nn.Module:
         return AttentionLayer(
             spatial_dims=self.dimensions,
             in_channels=in_channels,
@@ -249,6 +266,8 @@ def _get_bottom_layer(self, in_channels: int, out_channels: int, strides: int, l
                 strides=strides,
                 dropout=self.dropout,
             ),
+            up_kernel_size=self.up_kernel_size,
+            strides=strides,
             dropout=self.dropout,
         )
 

monai/transforms/__init__.py

@@ -227,6 +227,8 @@
 from .inverse_batch_transform import BatchInverseTransform, Decollated, DecollateD, DecollateDict
 from .io.array import SUPPORTED_READERS, LoadImage, SaveImage
 from .io.dictionary import LoadImaged, LoadImageD, LoadImageDict, SaveImaged, SaveImageD, SaveImageDict
+from .lazy.array import Apply
+from .lazy.functional import apply
 from .meta_utility.dictionary import (
     FromMetaTensord,
     FromMetaTensorD,
@@ -235,6 +237,13 @@
     ToMetaTensorD,
     ToMetaTensorDict,
 )
+from .meta_matrix import (
+    Grid,
+    matmul,
+    Matrix,
+    MatrixFactory,
+    MetaMatrix,
+)
 from .nvtx import (
     Mark,
     Markd,
@@ -449,7 +458,18 @@
     ZoomD,
     ZoomDict,
 )
-from .transform import MapTransform, Randomizable, RandomizableTransform, ThreadUnsafe, Transform, apply_transform
+from .transform import (
+    LazyTrait,
+    LazyTransform,
+    MapTransform,
+    MultiSampleTrait,
+    Randomizable,
+    RandomizableTrait,
+    RandomizableTransform,
+    ThreadUnsafe,
+    Transform,
+    apply_transform,
+)
 from .utility.array import (
     AddChannel,
     AddCoordinateChannels,
@@ -621,6 +641,8 @@
     generate_label_classes_crop_centers,
     generate_pos_neg_label_crop_centers,
     generate_spatial_bounding_box,
+    get_backend_from_tensor_like,
+    get_device_from_tensor_like,
     get_extreme_points,
     get_largest_connected_component_mask,
     get_number_image_type_conversions,

monai/transforms/lazy/__init__.py

@@ -0,0 +1,10 @@
+# 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.

monai/transforms/lazy/array.py

@@ -0,0 +1,48 @@
+# 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 monai.transforms.lazy.functional import apply
+from monai.transforms.inverse import InvertibleTransform
+
+__all__ = ["Apply"]
+
+
+class Apply(InvertibleTransform):
+    """
+    Apply wraps the apply method and can function as a Transform in either array or dictionary
+    mode.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, *args, **kwargs):
+        return apply(*args, **kwargs)
+
+    def inverse(self, data):
+        return NotImplementedError()
+
+
+# class Applyd(MapTransform, InvertibleTransform):
+#
+#     def __init__(self):
+#         super().__init__()
+#
+#     def __call__(
+#             self,
+#             d: dict
+#     ):
+#         rd = dict()
+#         for k, v in d.items():
+#             rd[k] = apply(v)
+#
+#     def inverse(self, data):
+#         return NotImplementedError()