Compose compile; initial multisample generic croppad (array and dict) implementations

Author: atbenmurray

monai/transforms/atmostonce/apply.py

@@ -110,6 +110,14 @@ def matrix_from_matrix_container(matrix):
 
 def apply(data: Union[torch.Tensor, MetaTensor],
           pending: Optional[dict] = None):
+
+    if isinstance(data, dict):
+        rd = dict()
+        for k, v in data.items():
+            result = apply(v)
+            rd[k] = result
+        return rd
+
     pending_ = pending
     pending_ = data.pending_transforms
 
@@ -188,17 +196,36 @@ def apply(data: Union[torch.Tensor, MetaTensor],
     return data
 
 
+# make Apply universal for arrays and dictionaries; it just calls through to functional apply
 class Apply(InvertibleTransform):
 
     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()
+
+
 # class Applyd(MapTransform, InvertibleTransform):
 #
 #     def __init__(self,

monai/transforms/atmostonce/array.py

@@ -737,15 +737,20 @@ def __call__(
             img: torch.Tensor,
             randomize: Optional[bool] = True
     ):
+
+        img_shape = img.shape[:1]
+
         if randomize:
             self.randomize(img)
 
         if self._do_transform:
             offsets_ = self.offsets
-            slices = tuple(slice(o, o + s) for o, s in zip(offsets_, self.sizes))
-            return self.op(img, slices=slices)
         else:
-            return self.op(img)
+            # center crop if this sample isn't random
+            offsets_ = tuple((i - s) // 2 for i, s in zip(img_shape, self.sizes))
+        slices = tuple(slice(o, o + s) for o, s in zip(offsets_, self.sizes))
+        return self.op(img, slices=slices)
+
 
     def inverse(
             self,

monai/transforms/atmostonce/compose.py

@@ -3,85 +3,86 @@
 
 import numpy as np
 
-
+from monai.transforms.atmostonce.apply import Apply
 from monai.transforms.atmostonce.lazy_transform import LazyTransform, compile_lazy_transforms, flatten_sequences
-from monai.transforms.atmostonce.utility import CachedTransformCompose
+from monai.transforms.atmostonce.utility import CachedTransformCompose, MultiSampleTransformCompose, \
+    IMultiSampleTransform, IRandomizableTransform, ILazyTransform
 from monai.utils import GridSampleMode, GridSamplePadMode, ensure_tuple, get_seed, MAX_SEED
 
 from monai.transforms import Randomizable, InvertibleTransform, OneOf, apply_transform
 
 
 # TODO: this is intended to replace Compose once development is done
 
-class ComposeCompiler:
-    """
-    Args:
-        transforms: A sequence of callable transforms
-        lazy_resampling: Whether to resample the data after each transform or accumulate
-            changes and then resample according to the accumulated changes as few times as
-            possible. Defaults to True as this nearly always improves speed and quality
-        caching_policy: Whether to cache deterministic transforms before the first
-            randomised transforms. This can be one of "off", "drive", "memory"
-        caching_favor: Whether to cache primarily for "speed" or for "quality". "speed" will
-            favor doing more work before caching, whereas "quality" will favour delaying
-            resampling until after caching
-    """
-    def __init__(
-            self,
-            transforms: Union[Sequence[Callable], Callable],
-            lazy_resampling: Optional[bool] = True,
-            caching_policy: Optional[str] = "off",
-            caching_favor: Optional[str] = "quality"
-    ):
-        valid_caching_policies = ("off", "drive", "memory")
-        if caching_policy not in valid_caching_policies:
-            raise ValueError("parameter 'caching_policy' must be one of "
-                             f"{valid_caching_policies} but is '{caching_policy}'")
-
-        dest_transforms = None
-
-        if caching_policy == "off":
-            if lazy_resampling is False:
-                dest_transforms = [t for t in transforms]
-            else:
-                dest_transforms = ComposeCompiler.lazy_no_cache()
-        else:
-            if caching_policy == "drive":
-                raise NotImplementedError()
-            elif caching_policy == "memory":
-                raise NotImplementedError()
-
-        self.src_transforms = [t for t in transforms]
-        self.dest_transforms = dest_transforms
-
-    def __getitem__(
-            self,
-            index
-    ):
-        return self.dest_transforms[index]
-
-    def __len__(self):
-        return len(self.dest_transforms)
-
-    @staticmethod
-    def lazy_no_cache(transforms):
-        dest_transforms = []
-        # TODO: replace with lazy transform
-        cur_lazy = []
-        for i_t in range(1, len(transforms)):
-            if isinstance(transforms[i_t], LazyTransform):
-                # add this to the stack of transforms to be handled lazily
-                    cur_lazy.append(transforms[i_t])
-            else:
-                if len(cur_lazy) > 0:
-                    dest_transforms.append(cur_lazy)
-                    # TODO: replace with lazy transform
-                    cur_lazy = []
-                dest_transforms.append(transforms[i_t])
-        return dest_transforms
+# class ComposeCompiler:
+#     """
+#     Args:
+#         transforms: A sequence of callable transforms
+#         lazy_resampling: Whether to resample the data after each transform or accumulate
+#             changes and then resample according to the accumulated changes as few times as
+#             possible. Defaults to True as this nearly always improves speed and quality
+#         caching_policy: Whether to cache deterministic transforms before the first
+#             randomised transforms. This can be one of "off", "drive", "memory"
+#         caching_favor: Whether to cache primarily for "speed" or for "quality". "speed" will
+#             favor doing more work before caching, whereas "quality" will favour delaying
+#             resampling until after caching
+#     """
+#     def __init__(
+#             self,
+#             transforms: Union[Sequence[Callable], Callable],
+#             lazy_resampling: Optional[bool] = True,
+#             caching_policy: Optional[str] = "off",
+#             caching_favor: Optional[str] = "quality"
+#     ):
+#         valid_caching_policies = ("off", "drive", "memory")
+#         if caching_policy not in valid_caching_policies:
+#             raise ValueError("parameter 'caching_policy' must be one of "
+#                              f"{valid_caching_policies} but is '{caching_policy}'")
+#
+#         dest_transforms = None
+#
+#         if caching_policy == "off":
+#             if lazy_resampling is False:
+#                 dest_transforms = [t for t in transforms]
+#             else:
+#                 dest_transforms = ComposeCompiler.lazy_no_cache()
+#         else:
+#             if caching_policy == "drive":
+#                 raise NotImplementedError()
+#             elif caching_policy == "memory":
+#                 raise NotImplementedError()
+#
+#         self.src_transforms = [t for t in transforms]
+#         self.dest_transforms = dest_transforms
+#
+#     def __getitem__(
+#             self,
+#             index
+#     ):
+#         return self.dest_transforms[index]
+#
+#     def __len__(self):
+#         return len(self.dest_transforms)
+#
+#     @staticmethod
+#     def lazy_no_cache(transforms):
+#         dest_transforms = []
+#         # TODO: replace with lazy transform
+#         cur_lazy = []
+#         for i_t in range(1, len(transforms)):
+#             if isinstance(transforms[i_t], LazyTransform):
+#                 # add this to the stack of transforms to be handled lazily
+#                     cur_lazy.append(transforms[i_t])
+#             else:
+#                 if len(cur_lazy) > 0:
+#                     dest_transforms.append(cur_lazy)
+#                     # TODO: replace with lazy transform
+#                     cur_lazy = []
+#                 dest_transforms.append(transforms[i_t])
+#         return dest_transforms
 
 
-class ComposeCompiler2:
+class ComposeCompiler:
 
     def compile(self, transforms, cache_mechanism):
 
@@ -93,26 +94,59 @@ def compile(self, transforms, cache_mechanism):
 
         return transforms___
 
-    def compile_caching(self, transforms, cache_stategy):
+    def compile_caching(self, transforms, cache_mechanism):
+        # TODO: handle being passed a transform list with containers
         # given a list of transforms, determine where to add a cached transform object
         # and what transforms to put in it
-        return transforms
+        cacheable = list()
+        for t in transforms:
+            if self.transform_is_random(t) is False:
+                cacheable.append(t)
+            else:
+                break
+
+        if len(cacheable) == 0:
+            return list(transforms)
+        else:
+            return [CachedTransformCompose(cacheable, cache_mechanism)] + transforms[len(cacheable):]
 
     def compile_multisampling(self, transforms):
-        return transforms
+        for i in reversed(range(len(transforms))):
+            if self.transform_is_multisampling(transforms[i]) is True:
+                transforms_ = transforms[:i] + [MultiSampleTransformCompose(transforms[i],
+                                                                           transforms[i+1:])]
+                return self.compile_multisampling(transforms_)
+
+        return list(transforms)
 
     def compile_lazy_resampling(self, transforms):
-        return transforms
+        result = list()
+        lazy = list()
+        for i in range(len(transforms)):
+            if self.transform_is_lazy(transforms[i]):
+                lazy.append(transforms[i])
+            else:
+                if len(lazy) > 0:
+                    result.extend(lazy)
+                    result.append(Apply())
+                    lazy = list()
+                result.append(transforms[i])
+        if len(lazy) > 0:
+            result.extend(lazy)
+            result.append(Apply())
+        return result
+
+    def transform_is_random(self, t):
+        return isinstance(t, IRandomizableTransform)
 
     def transform_is_container(self, t):
-        if isinstance(t, CachedTransform):
-            return True
-        return False
+        return isinstance(t, CachedTransformCompose, MultiSampleTransformCompose)
 
     def transform_is_multisampling(self, t):
-        # if isinstance(t, MultiSamplingTransform):
-        #     return True
-        return False
+        return isinstance(t, IMultiSampleTransform)
+
+    def transform_is_lazy(self, t):
+        return isinstance(t, ILazyTransform)
 
 
 class Compose(Randomizable, InvertibleTransform):