Assign group IDs to wrapped transforms in Compose

Extend Compose group tagging to discover traceable transforms wrapped in object attributes and container fields, including common patterns like transform/transforms wrappers.

Preserve nested Compose boundaries by setting group IDs only on the nested Compose instance itself, without traversing its internal pipeline.

Add a regression test covering wrapped transform attributes and nested container traversal.

Signed-off-by: sewon jeon <irocks0922@gmail.com>

Author: eclipse0922

monai/transforms/compose.py

@@ -277,39 +277,67 @@ def _set_transform_groups(self):
         group_id = str(id(self))
         visited = set()  # Track visited objects to avoid infinite recursion
 
-        def set_group_recursive(obj, gid):
+        def set_group_recursive(obj, gid, allow_compose: bool = False):
             """
             Recursively set a group ID on a transform and its wrapped transforms.
 
             Args:
                 obj: Transform instance to process.
                 gid: Group identifier to assign.
+                allow_compose: Whether to set group on ``Compose`` instances.
+                    ``Compose`` internals are not traversed to preserve nested
+                    pipeline boundaries.
 
             Returns:
                 None.
             """
+            if obj is None or isinstance(obj, (bool, int, float, str, bytes)):
+                return
+
             # Avoid infinite recursion
             obj_id = id(obj)
             if obj_id in visited:
                 return
             visited.add(obj_id)
 
+            if isinstance(obj, Compose):
+                if allow_compose:
+                    obj._group = gid
+                return
+
             if isinstance(obj, TraceableTransform):
                 obj._group = gid
 
-            # Handle wrapped transforms in dictionary transforms
-            # Check common attribute patterns for wrapped transforms
-            for attr_name in dir(obj):
-                # Skip magic methods and common non-transform attributes
-                if attr_name.startswith("__") or attr_name in ("transforms", "transform"):
-                    continue
-                attr = getattr(obj, attr_name, None)
-                if attr is not None and isinstance(attr, TraceableTransform) and not isinstance(attr, Compose):
-                    # Recursively set group on nested transforms
+            if isinstance(obj, Mapping):
+                for attr in obj.values():
+                    set_group_recursive(attr, gid)
+                return
+
+            if isinstance(obj, (list, tuple, set)):
+                for attr in obj:
                     set_group_recursive(attr, gid)
+                return
+
+            attrs: list[Any] = []
+            if hasattr(obj, "__dict__"):
+                attrs.extend(vars(obj).values())
+
+            slots = getattr(type(obj), "__slots__", ())
+            if isinstance(slots, str):
+                slots = (slots,)
+            for slot in slots:
+                if slot.startswith("__"):
+                    continue
+                try:
+                    attrs.append(getattr(obj, slot))
+                except AttributeError:
+                    continue
+
+            for attr in attrs:
+                set_group_recursive(attr, gid)
 
         for transform in self.transforms:
-            set_group_recursive(transform, group_id)
+            set_group_recursive(transform, group_id, allow_compose=True)
 
     @LazyTransform.lazy.setter  # type: ignore
     def lazy(self, val: bool):

tests/transforms/compose/test_compose.py

@@ -268,6 +268,43 @@ def test_data_loader_2(self):
             self.assertAlmostEqual(out_1.cpu().item(), 0.28602141572)
         set_determinism(None)
 
+    def test_set_transform_groups_on_wrapped_transform_attributes(self):
+        class _IdentityInvertible(mt.InvertibleTransform):
+            def __call__(self, data):
+                return data
+
+            def inverse(self, data):
+                return data
+
+        class _WrapperWithTransform:
+            def __init__(self):
+                self.transform = _IdentityInvertible()
+
+            def __call__(self, data):
+                return self.transform(data)
+
+        class _WrapperWithTransforms:
+            def __init__(self):
+                self.transforms = [_IdentityInvertible(), {"inner": _IdentityInvertible()}]
+
+            def __call__(self, data):
+                for transform in self.transforms:
+                    if isinstance(transform, dict):
+                        for nested_transform in transform.values():
+                            data = nested_transform(data)
+                    else:
+                        data = transform(data)
+                return data
+
+        wrapped_transform = _WrapperWithTransform()
+        wrapped_transforms = _WrapperWithTransforms()
+        composed = mt.Compose([wrapped_transform, wrapped_transforms])
+        expected_group = str(id(composed))
+
+        self.assertEqual(getattr(wrapped_transform.transform, "_group", None), expected_group)
+        self.assertEqual(getattr(wrapped_transforms.transforms[0], "_group", None), expected_group)
+        self.assertEqual(getattr(wrapped_transforms.transforms[1]["inner"], "_group", None), expected_group)
+
     def test_flatten_and_len(self):
         x = mt.EnsureChannelFirst(channel_dim="no_channel")
         t1 = mt.Compose([x, x, x, x, mt.Compose([mt.Compose([x, x]), x, x])])