allow map_items to be int

Signed-off-by: YunLiu <[email protected]>

Author: KumoLiu

monai/transforms/compose.py

@@ -47,7 +47,7 @@
 def execute_compose(
     data: NdarrayOrTensor | Sequence[NdarrayOrTensor] | Mapping[Any, NdarrayOrTensor],
     transforms: Sequence[Any],
-    map_items: bool = True,
+    map_items: bool | int = True,
     unpack_items: bool = False,
     start: int = 0,
     end: int | None = None,
@@ -66,7 +66,7 @@ def execute_compose(
         data: a tensor-like object to be transformed
         transforms: a sequence of transforms to be carried out
         map_items: whether to apply transform to each item in the input `data` if `data` is a list or tuple.
-            defaults to `True`.
+            defaults to `True`. If set to an integer, the transform will be applied to that index of the input `data`.
         unpack_items: whether to unpack input `data` with `*` as parameters for the callable function of transform.
             defaults to `False`.
         start: the index of the first transform to be executed. If not set, this defaults to 0
@@ -206,7 +206,7 @@ class Compose(Randomizable, InvertibleTransform, LazyTransform):
     Args:
         transforms: sequence of callables.
         map_items: whether to apply transform to each item in the input `data` if `data` is a list or tuple.
-            defaults to `True`.
+            defaults to `True`. If set to an integer, the transform will be applied to that index of the input `data`.
         unpack_items: whether to unpack input `data` with `*` as parameters for the callable function of transform.
             defaults to `False`.
         log_stats: this optional parameter allows you to specify a logger by name for logging of pipeline execution.
@@ -227,7 +227,7 @@ class Compose(Randomizable, InvertibleTransform, LazyTransform):
     def __init__(
         self,
         transforms: Sequence[Callable] | Callable | None = None,
-        map_items: bool = True,
+        map_items: bool | int = True,
         unpack_items: bool = False,
         log_stats: bool | str = False,
         lazy: bool | None = False,
@@ -238,9 +238,9 @@ def __init__(
         if transforms is None:
             transforms = []
 
-        if not isinstance(map_items, bool):
+        if not isinstance(map_items, (bool, int)):
             raise ValueError(
-                f"Argument 'map_items' should be boolean. Got {type(map_items)}."
+                f"Argument 'map_items' should be boolean or int. Got {type(map_items)}."
                 "Check brackets when passing a sequence of callables."
             )
 

monai/transforms/transform.py

@@ -101,7 +101,7 @@ def _apply_transform(
 def apply_transform(
     transform: Callable[..., ReturnType],
     data: Any,
-    map_items: bool = True,
+    map_items: bool | int = True,
     unpack_items: bool = False,
     log_stats: bool | str = False,
     lazy: bool | None = None,
@@ -119,6 +119,7 @@ def apply_transform(
         data: an object to be transformed.
         map_items: whether to apply transform to each item in `data`,
             if `data` is a list or tuple. Defaults to True.
+            it can also be an int, if so, apply the transform to each item in the list `map_items` times.
         unpack_items: whether to unpack parameters using `*`. Defaults to False.
         log_stats: log errors when they occur in the processing pipeline. By default, this is set to False, which
             disables the logger for processing pipeline errors. Setting it to None or True will enable logging to the
@@ -136,8 +137,15 @@ def apply_transform(
         Union[List[ReturnType], ReturnType]: The return type of `transform` or a list thereof.
     """
     try:
-        if isinstance(data, (list, tuple)) and map_items:
-            return [apply_transform(transform, item, map_items, unpack_items, log_stats, lazy, overrides) for item in data]
+        if isinstance(data, (list, tuple)) and isinstance(map_items, (int, bool)):
+            # if map_items is an int, apply the transform to each item in the list `map_items` times
+            if isinstance(map_items, int) and type(map_items) is not bool and map_items > 0:
+                return [
+                    apply_transform(transform, item, map_items - 2, unpack_items, log_stats, lazy, overrides)
+                    for item in data
+                ]
+            else:
+                return [_apply_transform(transform, item, unpack_items, lazy, overrides, log_stats) 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

tests/test_compose.py

@@ -141,6 +141,20 @@ def b(i, i2):
         self.assertEqual(mt.Compose(transforms, unpack_items=True)(data), expected)
         self.assertEqual(execute_compose(data, transforms, unpack_items=True), expected)
 
+    def test_list_non_dict_compose_with_unpack_map_2(self):
+
+        def a(i, i2):
+            return i + "a", i2 + "a2"
+
+        def b(i, i2):
+            return i + "b", i2 + "b2"
+
+        transforms = [a, b, a, b]
+        data = [[("", ""), ("", "")], [("t", "t"), ("t", "t")]]
+        expected = [[("abab", "a2b2a2b2"), ("abab", "a2b2a2b2")], [("tabab", "ta2b2a2b2"), ("tabab", "ta2b2a2b2")]]
+        self.assertEqual(mt.Compose(transforms, map_items=2, unpack_items=True)(data), expected)
+        self.assertEqual(execute_compose(data, transforms, map_items=2, unpack_items=True), expected)
+
     def test_list_dict_compose_no_map(self):
 
         def a(d):  # transform to handle dict data