changed compose to explicit flattening

Signed-off-by: Lukas Folle <[email protected]>

Author: lukas-folle-snkeos

monai/transforms/__init__.py

@@ -546,6 +546,7 @@
     TorchVision,
     ToTensor,
     Transpose,
+    FlattenSequence,
 )
 from .utility.dictionary import (
     AddCoordinateChannelsd,
@@ -671,6 +672,9 @@
     Transposed,
     TransposeD,
     TransposeDict,
+    FlattenSequenced,
+    FlattenSequenceD,
+    FlattenSequenceDict,
 )
 from .utils import (
     Fourier,

monai/transforms/traits.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-__all__ = ["LazyTrait", "InvertibleTrait", "RandomizableTrait", "MultiSampleTrait", "ThreadUnsafe"]
+__all__ = ["LazyTrait", "InvertibleTrait", "RandomizableTrait", "MultiSampleTrait", "ThreadUnsafe", "ReduceTrait"]
 
 from typing import Any
 
@@ -99,3 +99,14 @@ class ThreadUnsafe:
     """
 
     pass
+
+
+class ReduceTrait:
+    """
+    An interface to indicate that the transform has the capability to reduce multiple samples
+    into a single sample.
+    This interface can be extended from by people adapting transforms to the MONAI framework as well
+    as by implementors of MONAI transforms.
+    """
+
+    pass

monai/transforms/transform.py

@@ -25,7 +25,7 @@
 from monai import config, transforms
 from monai.config import KeysCollection
 from monai.data.meta_tensor import MetaTensor
-from monai.transforms.traits import LazyTrait, RandomizableTrait, ThreadUnsafe
+from monai.transforms.traits import LazyTrait, RandomizableTrait, ThreadUnsafe, ReduceTrait
 from monai.utils import MAX_SEED, ensure_tuple, first
 from monai.utils.enums import TransformBackends
 from monai.utils.misc import MONAIEnvVars
@@ -142,17 +142,11 @@ def apply_transform(
     """
     try:
         map_items_ = int(map_items) if isinstance(map_items, bool) else map_items
-        if isinstance(data, (list, tuple)) and map_items_ > 0:
-            res: list[Any] = []
-            for item in data:
-                res_item = apply_transform(transform, item, map_items_ - 1, unpack_items, log_stats, lazy, overrides)
-                # Only extend if we're at the leaf level (map_items_ == 1) and the transform
-                # actually returned a list (not preserving nested structure)
-                if isinstance(res_item, list) and map_items_ == 1 and not isinstance(item, (list, tuple)):
-                    res.extend(res_item)
-                else:
-                    res.append(res_item)
-            return res
+        if isinstance(data, (list, tuple)) and map_items_ > 0 and not isinstance(transform, ReduceTrait):
+            return [
+                apply_transform(transform, item, map_items_ - 1, unpack_items, log_stats, lazy, overrides)
+                for item in data
+            ]
         return _apply_transform(transform, data, unpack_items, lazy, overrides, log_stats)
     except Exception as e:
         # if in debug mode, don't swallow exception so that the breakpoint

monai/transforms/utility/array.py

@@ -43,7 +43,7 @@
     median_filter,
 )
 from monai.transforms.inverse import InvertibleTransform, TraceableTransform
-from monai.transforms.traits import MultiSampleTrait
+from monai.transforms.traits import MultiSampleTrait, ReduceTrait
 from monai.transforms.transform import Randomizable, RandomizableTrait, RandomizableTransform, Transform
 from monai.transforms.utils import (
     apply_affine_to_points,
@@ -110,6 +110,7 @@
     "ImageFilter",
     "RandImageFilter",
     "ApplyTransformToPoints",
+    "FlattenSequence"
 ]
 
 
@@ -1950,3 +1951,16 @@ def inverse(self, data: torch.Tensor) -> torch.Tensor:
             data = inverse_transform(data, transform[TraceKeys.EXTRA_INFO]["image_affine"])
 
         return data
+
+
+class FlattenSequence(Transform, ReduceTrait):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, data):
+        if isinstance(data, (list, tuple)):
+            if len(data) == 0:
+                return data
+            if isinstance(data[0], (list, tuple)):
+                return [item for sublist in data for item in sublist]
+        return data

monai/transforms/utility/dictionary.py

@@ -30,7 +30,7 @@
 from monai.data.meta_tensor import MetaObj, MetaTensor
 from monai.data.utils import no_collation
 from monai.transforms.inverse import InvertibleTransform
-from monai.transforms.traits import MultiSampleTrait, RandomizableTrait
+from monai.transforms.traits import MultiSampleTrait, RandomizableTrait, ReduceTrait
 from monai.transforms.transform import MapTransform, Randomizable, RandomizableTransform
 from monai.transforms.utility.array import (
     AddCoordinateChannels,
@@ -64,6 +64,7 @@
     TorchVision,
     ToTensor,
     Transpose,
+    FlattenSequence
 )
 from monai.transforms.utils import extreme_points_to_image, get_extreme_points
 from monai.transforms.utils_pytorch_numpy_unification import concatenate
@@ -191,6 +192,9 @@
     "ApplyTransformToPointsd",
     "ApplyTransformToPointsD",
     "ApplyTransformToPointsDict",
+    "FlattenSequenced",
+    "FlattenSequenceD",
+    "FlattenSequenceDict"
 ]
 
 DEFAULT_POST_FIX = PostFix.meta()
@@ -1906,6 +1910,23 @@ def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> dict[Hashable, torch
         return d
 
 
+class FlattenSequenced(MapTransform, ReduceTrait):
+    def __init__(
+        self,
+        keys: KeysCollection,
+        allow_missing_keys: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(keys, allow_missing_keys)
+        self.flatten_sequence = FlattenSequence(**kwargs)
+
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
+        d = dict(data)
+        for key in self.key_iterator(d):
+            d[key] = self.flatten_sequence(d[key])
+        return d
+
+
 RandImageFilterD = RandImageFilterDict = RandImageFilterd
 ImageFilterD = ImageFilterDict = ImageFilterd
 IdentityD = IdentityDict = Identityd
@@ -1949,3 +1970,4 @@ def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> dict[Hashable, torch
 AddCoordinateChannelsD = AddCoordinateChannelsDict = AddCoordinateChannelsd
 FlattenSubKeysD = FlattenSubKeysDict = FlattenSubKeysd
 ApplyTransformToPointsD = ApplyTransformToPointsDict = ApplyTransformToPointsd
+FlattenSequenceD = FlattenSequenceDict = FlattenSequenced

tests/transforms/compose/test_compose.py

@@ -285,26 +285,30 @@ def test_backwards_compatible_imports(self):
     def test_list_extend_multi_sample_trait(self):
         center_crop = mt.CenterSpatialCrop([128, 128])
         multi_sample_transform = mt.RandSpatialCropSamples([64, 64], 1)
+        flatten_sequence_transform = mt.FlattenSequence()
 
         img = torch.zeros([1, 512, 512])
 
         self.assertEqual(execute_compose(img, [center_crop]).shape, torch.Size([1, 128, 128]))
-        single_multi_sample_trait_result = execute_compose(img, [multi_sample_transform, center_crop])
+        single_multi_sample_trait_result = execute_compose(img, [multi_sample_transform, center_crop, flatten_sequence_transform])
         self.assertIsInstance(single_multi_sample_trait_result, list)
         self.assertEqual(len(single_multi_sample_trait_result), 1)
         self.assertEqual(single_multi_sample_trait_result[0].shape, torch.Size([1, 64, 64]))
 
-        double_multi_sample_trait_result = execute_compose(img, [multi_sample_transform, multi_sample_transform, center_crop])
+        double_multi_sample_trait_result = execute_compose(img, [
+            multi_sample_transform, multi_sample_transform, flatten_sequence_transform, center_crop
+        ])
         self.assertIsInstance(double_multi_sample_trait_result, list)
         self.assertEqual(len(double_multi_sample_trait_result), 1)
         self.assertEqual(double_multi_sample_trait_result[0].shape, torch.Size([1, 64, 64]))
 
     def test_multi_sample_trait_cardinality(self):
         img = torch.zeros([1, 128, 128])
         t2 = mt.RandSpatialCropSamples([32, 32], num_samples=2)
+        flatten_sequence_transform = mt.FlattenSequence()
 
         # chaining should multiply counts: 2 x 2 = 4, flattened
-        res = execute_compose(img, [t2, t2])
+        res = execute_compose(img, [t2, t2, flatten_sequence_transform])
         self.assertIsInstance(res, list)
         self.assertEqual(len(res), 4)
         for r in res: