added test for many multisample transforms; refactored code

Author: lukas-folle-snkeos

monai/transforms/transform.py

@@ -148,11 +148,8 @@ def apply_transform(
                 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:
-                    if not isinstance(item, (list, tuple)):
-                        res.extend(res_item)
-                    else:
-                        res.append(res_item)
+                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

tests/transforms/compose/test_compose.py

@@ -282,6 +282,28 @@ def test_flatten_and_len(self):
     def test_backwards_compatible_imports(self):
         from monai.transforms.transform import MapTransform, RandomizableTransform, Transform  # noqa: F401
 
+    def test_list_extend_multi_sample_trait(self):
+        from monai.transforms import CenterSpatialCrop, RandSpatialCropSamples
+
+        center_crop = CenterSpatialCrop([128, 128])
+        multi_sample_transform = RandSpatialCropSamples([64, 64], 1)
+
+        img = torch.zeros([1, 512, 512])
+
+        assert execute_compose(img, [center_crop]).shape == torch.Size([1, 128, 128])
+        single_multi_sample_trait_result = execute_compose(img, [multi_sample_transform, center_crop])
+        assert (
+            isinstance(single_multi_sample_trait_result, list)
+            and len(single_multi_sample_trait_result) == 1
+            and 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])
+        assert (
+            isinstance(double_multi_sample_trait_result, list)
+            and len(double_multi_sample_trait_result) == 1
+            and double_multi_sample_trait_result[0].shape == torch.Size([1, 64, 64])
+        )
+
 
 TEST_COMPOSE_EXECUTE_TEST_CASES = [
     [None, tuple()],