Improve Compose encapsulation

monai/data/dataset.py

@@ -39,7 +39,6 @@
     Compose,
     Randomizable,
     RandomizableTrait,
-    ThreadUnsafe,
     Transform,
     apply_transform,
     convert_to_contiguous,
@@ -316,13 +315,15 @@ def _pre_transform(self, item_transformed):
             random transform object
 
         """
-        for _transform in self.transform.transforms:
-            # execute all the deterministic transforms
-            if isinstance(_transform, RandomizableTrait) or not isinstance(_transform, Transform):
-                break
-            # this is to be consistent with CacheDataset even though it's not in a multi-thread situation.
-            _xform = deepcopy(_transform) if isinstance(_transform, ThreadUnsafe) else _transform
-            item_transformed = apply_transform(_xform, item_transformed)
+        if not isinstance(self.transform, Compose):
+            raise ValueError("transform must be an instance of monai.transforms.Compose.")
+
+        first_random = self.transform.get_index_of_first(
+            lambda t: isinstance(t, RandomizableTrait) or not isinstance(t, Transform)
+        )
+
+        item_transformed = self.transform(item_transformed, end=first_random, threading=True)
+
         if self.reset_ops_id:
             reset_ops_id(item_transformed)
         return item_transformed
@@ -340,15 +341,12 @@ def _post_transform(self, item_transformed):
         """
         if not isinstance(self.transform, Compose):
             raise ValueError("transform must be an instance of monai.transforms.Compose.")
-        start_post_randomize_run = False
-        for _transform in self.transform.transforms:
-            if (
-                start_post_randomize_run
-                or isinstance(_transform, RandomizableTrait)
-                or not isinstance(_transform, Transform)
-            ):
-                start_post_randomize_run = True
-                item_transformed = apply_transform(_transform, item_transformed)
+
+        first_random = self.transform.get_index_of_first(
+            lambda t: isinstance(t, RandomizableTrait) or not isinstance(t, Transform)
+        )
+        if first_random is not None:
+            item_transformed = self.transform(item_transformed, start=first_random)
         return item_transformed
 
     def _cachecheck(self, item_transformed):
@@ -492,11 +490,9 @@ def _pre_transform(self, item_transformed):
         """
         if not isinstance(self.transform, Compose):
             raise ValueError("transform must be an instance of monai.transforms.Compose.")
-        for i, _transform in enumerate(self.transform.transforms):
-            if i == self.cache_n_trans:
-                break
-            _xform = deepcopy(_transform) if isinstance(_transform, ThreadUnsafe) else _transform
-            item_transformed = apply_transform(_xform, item_transformed)
+
+        item_transformed = self.transform(item_transformed, end=self.cache_n_trans, threading=True)
+
         reset_ops_id(item_transformed)
         return item_transformed
 
@@ -512,10 +508,8 @@ def _post_transform(self, item_transformed):
         """
         if not isinstance(self.transform, Compose):
             raise ValueError("transform must be an instance of monai.transforms.Compose.")
-        for i, _transform in enumerate(self.transform.transforms):
-            if i >= self.cache_n_trans:
-                item_transformed = apply_transform(_transform, item_transformed)
-        return item_transformed
+
+        return self.transform(item_transformed, start=self.cache_n_trans)
 
 
 class LMDBDataset(PersistentDataset):
@@ -879,12 +873,12 @@ def _load_cache_item(self, idx: int):
             idx: the index of the input data sequence.
         """
         item = self.data[idx]
-        for _transform in self.transform.transforms:
-            # execute all the deterministic transforms
-            if isinstance(_transform, RandomizableTrait) or not isinstance(_transform, Transform):
-                break
-            _xform = deepcopy(_transform) if isinstance(_transform, ThreadUnsafe) else _transform
-            item = apply_transform(_xform, item)
+
+        first_random = self.transform.get_index_of_first(
+            lambda t: isinstance(t, RandomizableTrait) or not isinstance(t, Transform)
+        )
+        item = self.transform(item, end=first_random, threading=True)
+
         if self.as_contiguous:
             item = convert_to_contiguous(item, memory_format=torch.contiguous_format)
         return item
@@ -911,17 +905,15 @@ def _transform(self, index: int):
             data = self._cache[cache_index] = self._load_cache_item(cache_index)
 
         # load data from cache and execute from the first random transform
-        start_run = False
         if not isinstance(self.transform, Compose):
             raise ValueError("transform must be an instance of monai.transforms.Compose.")
-        for _transform in self.transform.transforms:
-            if start_run or isinstance(_transform, RandomizableTrait) or not isinstance(_transform, Transform):
-                # only need to deep copy data on first non-deterministic transform
-                if not start_run:
-                    start_run = True
-                    if self.copy_cache:
-                        data = deepcopy(data)
-                data = apply_transform(_transform, data)
+
+        first_random = self.transform.get_index_of_first(
+            lambda t: isinstance(t, RandomizableTrait) or not isinstance(t, Transform)
+        )
+        if first_random is not None:
+            data = self.transform(deepcopy(data), start=first_random)
+
         return data
 
 

monai/transforms/compose.py

@@ -16,12 +16,15 @@
 
 import warnings
 from collections.abc import Callable, Mapping, Sequence
+from copy import deepcopy
 from typing import Any
 
 import numpy as np
 
 import monai
+from monai.config import NdarrayOrTensor
 from monai.transforms.inverse import InvertibleTransform
+from monai.transforms.traits import ThreadUnsafe
 
 # For backwards compatibility (so this still works: from monai.transforms.compose import MapTransform)
 from monai.transforms.transform import (  # noqa: F401
@@ -152,6 +155,12 @@ def randomize(self, data: Any | None = None) -> None:
                     f'Transform "{tfm_name}" in Compose not randomized\n{tfm_name}.{type_error}.', RuntimeWarning
                 )
 
+    def get_index_of_first(self, predicate):
+        for i in range(len(self.transforms)):
+            if predicate(self.transforms[i]):
+                return i
+        return None
+
     def flatten(self):
         """Return a Composition with a simple list of transforms, as opposed to any nested Compositions.
 
@@ -172,11 +181,69 @@ def __len__(self):
         """Return number of transformations."""
         return len(self.flatten().transforms)
 
-    def __call__(self, input_):
-        for _transform in self.transforms:
-            input_ = apply_transform(_transform, input_, self.map_items, self.unpack_items, self.log_stats)
+    @classmethod
+    def execute(
+        cls,
+        input_: NdarrayOrTensor,
+        transforms: Sequence[Any],
+        map_items: bool = True,
+        unpack_items: bool = False,
+        log_stats: bool = False,
+        start: int = 0,
+        end: int | None = None,
+        threading: bool = False,
+    ) -> NdarrayOrTensor:
+        """
+        ``execute`` provides the implementation that Compose uses to execute a sequence
+        of transforms. As well as being used by Compose, it can be used by subclasses of
+        Compose and by code that doesn't have a Compose instance but needs to execute a
+        sequence of transforms is if it were executed by Compose. For the most part, it
+        is recommended to use Compose instances, however.
+        Args:
+            `input_`: a tensor-like object to be transformed
+            transforms: a sequence of transforms to be carried out
+            map_items: whether to apply the transform to each item in ``data```.
+            Defaults to True if not set.
+            unpack_items: whether to unpack parameters using '*'. Defaults to False if not set
+            log_stats: whether to log detailed information about the application of ``transforms``
+            to ``input_``. For NumPy ndarrays and PyTorch tensors, log only the data shape and
+            value range. Defaults to False if not set.
+            start:
+            end:
+            threading:
+
+        Returns:
+
+        """
+        end_ = len(transforms) if end is None else end
+        if start is None:
+            raise ValueError(f"'start' ({start}) cannot be None")
+        if start > end_:
+            raise ValueError(f"'start' ({start}) must be less than 'end' ({end_})")
+        if end_ > len(transforms):
+            raise ValueError(f"'end' ({end_}) must be less than or equal to the transform count ({len(transforms)}")
+
+        # no-op if the range is empty
+        if start == end:
+            return input_
+
+        for _transform in transforms[start:end]:
+            if threading:
+                _transform = deepcopy(_transform) if isinstance(_transform, ThreadUnsafe) else _transform
+            input_ = apply_transform(_transform, input_, map_items, unpack_items, log_stats)  # type: ignore
         return input_
 
+    def __call__(self, input_, start=0, end=None, threading=False):
+        return Compose.execute(
+            input_,
+            self.transforms,
+            map_items=self.map_items,
+            unpack_items=self.unpack_items,
+            start=start,
+            end=end,
+            threading=threading,
+        )
+
     def inverse(self, data):
         invertible_transforms = [t for t in self.flatten().transforms if isinstance(t, InvertibleTransform)]
         if not invertible_transforms:
@@ -254,12 +321,23 @@ def flatten(self):
                 weights.append(w)
         return OneOf(transforms, weights, self.map_items, self.unpack_items)
 
-    def __call__(self, data):
+    def __call__(self, data, start=0, end=None, threading=False):
         if len(self.transforms) == 0:
             return data
+
         index = self.R.multinomial(1, self.weights).argmax()
         _transform = self.transforms[index]
-        data = apply_transform(_transform, data, self.map_items, self.unpack_items, self.log_stats)
+
+        data = Compose.execute(
+            data,
+            [_transform],
+            map_items=self.map_items,
+            unpack_items=self.unpack_items,
+            start=start,
+            end=end,
+            threading=threading,
+        )
+
         # if the data is a mapping (dictionary), append the OneOf transform to the end
         if isinstance(data, monai.data.MetaTensor):
             self.push_transform(data, extra_info={"index": index})
@@ -318,14 +396,22 @@ def __init__(
     ) -> None:
         super().__init__(transforms, map_items, unpack_items, log_stats)
 
-    def __call__(self, input_):
+    def __call__(self, input_, start=0, end=None, threading=False):
         if len(self.transforms) == 0:
             return input_
         num = len(self.transforms)
         applied_order = self.R.permutation(range(num))
 
-        for index in applied_order:
-            input_ = apply_transform(self.transforms[index], input_, self.map_items, self.unpack_items, self.log_stats)
+        input_ = Compose.execute(
+            input_,
+            [self.transforms[ind] for ind in applied_order],
+            map_items=self.map_items,
+            unpack_items=self.unpack_items,
+            start=start,
+            end=end,
+            threading=threading,
+        )
+
         # if the data is a mapping (dictionary), append the RandomOrder transform to the end
         if isinstance(input_, monai.data.MetaTensor):
             self.push_transform(input_, extra_info={"applied_order": applied_order})