nrtk-albumentations 2.1.0__py3-none-any.whl

albumentations/core/composition.py

@@ -0,0 +1,1885 @@
+"""Module for composing multiple transforms into augmentation pipelines.
+
+This module provides classes for combining multiple transformations into cohesive
+augmentation pipelines. It includes various composition strategies such as sequential
+application, random selection, and conditional application of transforms. These
+composition classes handle the coordination between different transforms, ensuring
+proper data flow and maintaining consistent behavior across the augmentation pipeline.
+"""
+
+from __future__ import annotations
+
+import random
+import warnings
+from collections import defaultdict
+from collections.abc import Iterator, Sequence
+from typing import Any, Union, cast
+
+import cv2
+import numpy as np
+
+from .bbox_utils import BboxParams, BboxProcessor
+from .hub_mixin import HubMixin
+from .keypoints_utils import KeypointParams, KeypointsProcessor
+from .serialization import (
+    SERIALIZABLE_REGISTRY,
+    Serializable,
+    get_shortest_class_fullname,
+    instantiate_nonserializable,
+)
+from .transforms_interface import BasicTransform
+from .utils import DataProcessor, format_args, get_shape
+
+__all__ = [
+    "BaseCompose",
+    "BboxParams",
+    "Compose",
+    "KeypointParams",
+    "OneOf",
+    "OneOrOther",
+    "RandomOrder",
+    "ReplayCompose",
+    "SelectiveChannelTransform",
+    "Sequential",
+    "SomeOf",
+]
+
+NUM_ONEOF_TRANSFORMS = 2
+REPR_INDENT_STEP = 2
+
+TransformType = Union[BasicTransform, "BaseCompose"]
+TransformsSeqType = list[TransformType]
+
+AVAILABLE_KEYS = ("image", "mask", "masks", "bboxes", "keypoints", "volume", "volumes", "mask3d", "masks3d")
+
+MASK_KEYS = (
+    "mask",  # 2D mask
+    "masks",  # Multiple 2D masks
+    "mask3d",  # 3D mask
+    "masks3d",  # Multiple 3D masks
+)
+
+# Keys related to image data
+IMAGE_KEYS = {"image", "images"}
+CHECKED_SINGLE = {"image", "mask"}
+CHECKED_MULTI = {"masks", "images", "volumes", "masks3d"}
+CHECK_BBOX_PARAM = {"bboxes"}
+CHECK_KEYPOINTS_PARAM = {"keypoints"}
+VOLUME_KEYS = {"volume", "volumes"}
+CHECKED_VOLUME = {"volume"}
+CHECKED_VOLUMES = {"volumes"}
+CHECKED_MASK3D = {"mask3d"}
+CHECKED_MASKS3D = {"masks3d"}
+
+
+class BaseCompose(Serializable):
+    """Base class for composing multiple transforms together.
+
+    This class serves as a foundation for creating compositions of transforms
+    in the Albumentations library. It provides basic functionality for
+    managing a sequence of transforms and applying them to data.
+
+    The class supports dynamic pipeline modification after initialization using
+    mathematical operators:
+    - Addition (`+`): Add transforms to the end of the pipeline
+    - Right addition (`__radd__`): Add transforms to the beginning of the pipeline
+    - Subtraction (`-`): Remove transforms by class from the pipeline
+
+    Attributes:
+        transforms (List[TransformType]): A list of transforms to be applied.
+        p (float): Probability of applying the compose. Should be in the range [0, 1].
+        replay_mode (bool): If True, the compose is in replay mode.
+        _additional_targets (Dict[str, str]): Additional targets for transforms.
+        _available_keys (Set[str]): Set of available keys for data.
+        processors (Dict[str, Union[BboxProcessor, KeypointsProcessor]]): Processors for specific data types.
+
+    Args:
+        transforms (TransformsSeqType): A sequence of transforms to compose.
+        p (float): Probability of applying the compose.
+
+    Raises:
+        ValueError: If an invalid additional target is specified.
+
+    Note:
+        - Subclasses should implement the __call__ method to define how
+          the composition is applied to data.
+        - The class supports serialization and deserialization of transforms.
+        - It provides methods for adding targets, setting deterministic behavior,
+          and checking data validity post-transform.
+        - All compose classes support pipeline modification operators:
+          - `compose + transform` adds individual transform(s) to the end
+          - `transform + compose` adds individual transform(s) to the beginning
+          - `compose - TransformClass` removes transforms by class type
+          - Only BasicTransform instances (not BaseCompose) can be added
+        - All operator operations return new instances without modifying the original.
+
+    Examples:
+        >>> import albumentations as A
+        >>> # Create base pipeline
+        >>> compose = A.Compose([A.HorizontalFlip(p=1.0)])
+        >>>
+        >>> # Add transforms using operators
+        >>> extended = compose + A.VerticalFlip(p=1.0)  # Append
+        >>> extended = compose + [A.Blur(), A.Rotate()]  # Append multiple
+        >>> extended = A.RandomCrop(256, 256) + compose  # Prepend
+        >>>
+        >>> # Remove transforms by class
+        >>> compose = A.Compose([A.HorizontalFlip(p=0.5), A.VerticalFlip(p=1.0)])
+        >>> reduced = compose - A.HorizontalFlip  # Remove by class
+
+    """
+
+    _transforms_dict: dict[int, BasicTransform] | None = None
+    check_each_transform: tuple[DataProcessor, ...] | None = None
+    main_compose: bool = True
+
+    def __init__(
+        self,
+        transforms: TransformsSeqType,
+        p: float,
+        mask_interpolation: int | None = None,
+        seed: int | None = None,
+        save_applied_params: bool = False,
+        **kwargs: Any,
+    ):
+        if isinstance(transforms, (BaseCompose, BasicTransform)):
+            warnings.warn(
+                "transforms is single transform, but a sequence is expected! Transform will be wrapped into list.",
+                stacklevel=2,
+            )
+            transforms = [transforms]
+
+        self.transforms = transforms
+        self.p = p
+
+        self.replay_mode = False
+        self._additional_targets: dict[str, str] = {}
+        self._available_keys: set[str] = set()
+        self.processors: dict[str, BboxProcessor | KeypointsProcessor] = {}
+        self._set_keys()
+        self.set_mask_interpolation(mask_interpolation)
+        self.set_random_seed(seed)
+        self.save_applied_params = save_applied_params
+
+    def _track_transform_params(self, transform: TransformType, data: dict[str, Any]) -> None:
+        """Track transform parameters if tracking is enabled."""
+        if "applied_transforms" in data and hasattr(transform, "params") and transform.params:
+            data["applied_transforms"].append((transform.__class__.__name__, transform.params.copy()))
+
+    def set_random_state(
+        self,
+        random_generator: np.random.Generator,
+        py_random: random.Random,
+    ) -> None:
+        """Set random state directly from generators.
+
+        Args:
+            random_generator (np.random.Generator): numpy random generator to use
+            py_random (random.Random): python random generator to use
+
+        """
+        self.random_generator = random_generator
+        self.py_random = py_random
+
+        # Propagate both random states to all transforms
+        for transform in self.transforms:
+            if isinstance(transform, (BasicTransform, BaseCompose)):
+                transform.set_random_state(random_generator, py_random)
+
+    def set_random_seed(self, seed: int | None) -> None:
+        """Set random state from seed.
+
+        Args:
+            seed (int | None): Random seed to use
+
+        """
+        # Store the original seed
+        self.seed = seed
+
+        # Use base seed directly (subclasses like Compose can override this)
+        self.random_generator = np.random.default_rng(seed)
+        self.py_random = random.Random(seed)
+
+        # Propagate seed to all transforms
+        for transform in self.transforms:
+            if isinstance(transform, (BasicTransform, BaseCompose)):
+                transform.set_random_seed(seed)
+
+    def set_mask_interpolation(self, mask_interpolation: int | None) -> None:
+        """Set interpolation mode for mask resizing operations.
+
+        Args:
+            mask_interpolation (int | None): OpenCV interpolation flag to use for mask transforms.
+                If None, default interpolation for masks will be used.
+
+        """
+        self.mask_interpolation = mask_interpolation
+        self._set_mask_interpolation_recursive(self.transforms)
+
+    def _set_mask_interpolation_recursive(self, transforms: TransformsSeqType) -> None:
+        for transform in transforms:
+            if isinstance(transform, BasicTransform):
+                if hasattr(transform, "mask_interpolation") and self.mask_interpolation is not None:
+                    transform.mask_interpolation = self.mask_interpolation
+            elif isinstance(transform, BaseCompose):
+                transform.set_mask_interpolation(self.mask_interpolation)
+
+    def __iter__(self) -> Iterator[TransformType]:
+        return iter(self.transforms)
+
+    def __len__(self) -> int:
+        return len(self.transforms)
+
+    def __call__(self, *args: Any, **data: Any) -> dict[str, Any]:
+        """Apply transforms.
+
+        Args:
+            *args (Any): Positional arguments are not supported.
+            **data (Any): Named parameters with data to transform.
+
+        Returns:
+            dict[str, Any]: Transformed data.
+
+        Raises:
+            NotImplementedError: This method must be implemented by subclasses.
+
+        """
+        raise NotImplementedError
+
+    def __getitem__(self, item: int) -> TransformType:
+        return self.transforms[item]
+
+    def __repr__(self) -> str:
+        return self.indented_repr()
+
+    @property
+    def additional_targets(self) -> dict[str, str]:
+        """Get additional targets dictionary.
+
+        Returns:
+            dict[str, str]: Dictionary containing additional targets mapping.
+
+        """
+        return self._additional_targets
+
+    @property
+    def available_keys(self) -> set[str]:
+        """Get set of available keys.
+
+        Returns:
+            set[str]: Set of string keys available for transforms.
+
+        """
+        return self._available_keys
+
+    def indented_repr(self, indent: int = REPR_INDENT_STEP) -> str:
+        """Get an indented string representation of the composition.
+
+        Args:
+            indent (int): Indentation level. Default: REPR_INDENT_STEP.
+
+        Returns:
+            str: Formatted string representation with proper indentation.
+
+        """
+        args = {k: v for k, v in self.to_dict_private().items() if not (k.startswith("__") or k == "transforms")}
+        repr_string = self.__class__.__name__ + "(["
+        for t in self.transforms:
+            repr_string += "\n"
+            t_repr = t.indented_repr(indent + REPR_INDENT_STEP) if hasattr(t, "indented_repr") else repr(t)
+            repr_string += " " * indent + t_repr + ","
+        repr_string += "\n" + " " * (indent - REPR_INDENT_STEP) + f"], {format_args(args)})"
+        return repr_string
+
+    @classmethod
+    def get_class_fullname(cls) -> str:
+        """Get the full qualified name of the class.
+
+        Returns:
+            str: The shortest class fullname.
+
+        """
+        return get_shortest_class_fullname(cls)
+
+    @classmethod
+    def is_serializable(cls) -> bool:
+        """Check if the class is serializable.
+
+        Returns:
+            bool: True if the class is serializable, False otherwise.
+
+        """
+        return True
+
+    def to_dict_private(self) -> dict[str, Any]:
+        """Convert the composition to a dictionary for serialization.
+
+        Returns:
+            dict[str, Any]: Dictionary representation of the composition.
+
+        """
+        return {
+            "__class_fullname__": self.get_class_fullname(),
+            "p": self.p,
+            "transforms": [t.to_dict_private() for t in self.transforms],
+        }
+
+    def get_dict_with_id(self) -> dict[str, Any]:
+        """Get a dictionary representation with object IDs for replay mode.
+
+        Returns:
+            dict[str, Any]: Dictionary with composition data and object IDs.
+
+        """
+        return {
+            "__class_fullname__": self.get_class_fullname(),
+            "id": id(self),
+            "params": None,
+            "transforms": [t.get_dict_with_id() for t in self.transforms],
+        }
+
+    def add_targets(self, additional_targets: dict[str, str] | None) -> None:
+        """Add additional targets to all transforms.
+
+        Args:
+            additional_targets (dict[str, str] | None): Dict of name -> type mapping for additional targets.
+                If None, no additional targets will be added.
+
+        """
+        if additional_targets:
+            for k, v in additional_targets.items():
+                if k in self._additional_targets and v != self._additional_targets[k]:
+                    raise ValueError(
+                        f"Trying to overwrite existed additional targets. "
+                        f"Key={k} Exists={self._additional_targets[k]} New value: {v}",
+                    )
+            self._additional_targets.update(additional_targets)
+            for t in self.transforms:
+                t.add_targets(additional_targets)
+            for proc in self.processors.values():
+                proc.add_targets(additional_targets)
+        self._set_keys()
+
+    def _set_keys(self) -> None:
+        """Set _available_keys"""
+        self._available_keys.update(self._additional_targets.keys())
+        for t in self.transforms:
+            self._available_keys.update(t.available_keys)
+            if hasattr(t, "targets_as_params"):
+                self._available_keys.update(t.targets_as_params)
+        if self.processors:
+            self._available_keys.update(["labels"])
+            for proc in self.processors.values():
+                if proc.default_data_name not in self._available_keys:  # if no transform to process this data
+                    warnings.warn(
+                        f"Got processor for {proc.default_data_name}, but no transform to process it.",
+                        stacklevel=2,
+                    )
+                self._available_keys.update(proc.data_fields)
+                if proc.params.label_fields:
+                    self._available_keys.update(proc.params.label_fields)
+
+    def set_deterministic(self, flag: bool, save_key: str = "replay") -> None:
+        """Set deterministic mode for all transforms.
+
+        Args:
+            flag (bool): Whether to enable deterministic mode.
+            save_key (str): Key to save replay parameters. Default: "replay".
+
+        """
+        for t in self.transforms:
+            t.set_deterministic(flag, save_key)
+
+    def check_data_post_transform(self, data: dict[str, Any]) -> dict[str, Any]:
+        """Check and filter data after transformation.
+
+        Args:
+            data (dict[str, Any]): Dictionary containing transformed data
+
+        Returns:
+            dict[str, Any]: Filtered data dictionary
+
+        """
+        if self.check_each_transform:
+            shape = get_shape(data)
+
+            for proc in self.check_each_transform:
+                for data_name, data_value in data.items():
+                    if data_name in proc.data_fields or (
+                        data_name in self._additional_targets
+                        and self._additional_targets[data_name] in proc.data_fields
+                    ):
+                        data[data_name] = proc.filter(data_value, shape)
+        return data
+
+    def _validate_transforms(self, transforms: list[Any]) -> None:
+        """Validate that all elements are BasicTransform instances.
+
+        Args:
+            transforms: List of objects to validate
+
+        Raises:
+            TypeError: If any element is not a BasicTransform instance
+
+        """
+        for t in transforms:
+            if not isinstance(t, BasicTransform):
+                raise TypeError(
+                    f"All elements must be instances of BasicTransform, got {type(t).__name__}",
+                )
+
+    def _combine_transforms(self, other: TransformType | TransformsSeqType, *, prepend: bool = False) -> BaseCompose:
+        """Combine transforms with the current compose.
+
+        Args:
+            other: Transform or sequence of transforms to combine
+            prepend: If True, prepend other to the beginning; if False, append to the end
+
+        Returns:
+            BaseCompose: New compose instance with combined transforms
+
+        Raises:
+            TypeError: If other is not a valid transform or sequence of transforms
+
+        """
+        if isinstance(other, (list, tuple)):
+            self._validate_transforms(other)
+            other_list = list(other)
+        else:
+            self._validate_transforms([other])
+            other_list = [other]
+
+        new_transforms = [*other_list, *list(self.transforms)] if prepend else [*list(self.transforms), *other_list]
+
+        return self._create_new_instance(new_transforms)
+
+    def __add__(self, other: TransformType | TransformsSeqType) -> BaseCompose:
+        """Add transform(s) to the end of this compose.
+
+        Args:
+            other: Transform or sequence of transforms to append
+
+        Returns:
+            BaseCompose: New compose instance with transforms appended
+
+        Raises:
+            TypeError: If other is not a valid transform or sequence of transforms
+
+        Examples:
+            >>> new_compose = compose + A.HorizontalFlip()
+            >>> new_compose = compose + [A.HorizontalFlip(), A.VerticalFlip()]
+
+        """
+        return self._combine_transforms(other, prepend=False)
+
+    def __radd__(self, other: TransformType | TransformsSeqType) -> BaseCompose:
+        """Add transform(s) to the beginning of this compose.
+
+        Args:
+            other: Transform or sequence of transforms to prepend
+
+        Returns:
+            BaseCompose: New compose instance with transforms prepended
+
+        Raises:
+            TypeError: If other is not a valid transform or sequence of transforms
+
+        Examples:
+            >>> new_compose = A.HorizontalFlip() + compose
+            >>> new_compose = [A.HorizontalFlip(), A.VerticalFlip()] + compose
+
+        """
+        return self._combine_transforms(other, prepend=True)
+
+    def __sub__(self, other: type[BasicTransform]) -> BaseCompose:
+        """Remove transform from this compose by class type.
+
+        Removes the first transform in the compose that matches the provided transform class.
+
+        Args:
+            other: Transform class to remove (e.g., A.HorizontalFlip)
+
+        Returns:
+            BaseCompose: New compose instance with transform removed
+
+        Raises:
+            TypeError: If other is not a BasicTransform class
+            ValueError: If no transform of that type is found in the compose
+
+        Note:
+            If multiple transforms of the same type exist in the compose,
+            only the first occurrence will be removed.
+
+        Examples:
+            >>> # Remove by transform class
+            >>> new_compose = compose - A.HorizontalFlip
+            >>>
+            >>> # With duplicates - only first occurrence removed
+            >>> compose = A.Compose([A.HorizontalFlip(p=0.5), A.VerticalFlip(), A.HorizontalFlip(p=1.0)])
+            >>> result = compose - A.HorizontalFlip  # Removes first HorizontalFlip (p=0.5)
+            >>> len(result.transforms)  # 2 (VerticalFlip and second HorizontalFlip remain)
+
+        """
+        # Validate that other is a BasicTransform class
+        if not (isinstance(other, type) and issubclass(other, BasicTransform)):
+            raise TypeError(
+                f"Can only remove BasicTransform classes, got {type(other).__name__}",
+            )
+
+        # Find first transform of matching class
+        new_transforms = list(self.transforms)
+        for i, transform in enumerate(new_transforms):
+            if type(transform) is other:
+                new_transforms.pop(i)
+                return self._create_new_instance(new_transforms)
+
+        # No matching transform found
+        class_name = other.__name__
+        raise ValueError(f"No transform of type {class_name} found in the compose pipeline")
+
+    def _create_new_instance(self, new_transforms: TransformsSeqType) -> BaseCompose:
+        """Create a new instance of the same class with new transforms.
+
+        Args:
+            new_transforms: List of transforms for the new instance
+
+        Returns:
+            BaseCompose: New instance of the same class
+
+        """
+        # Get current instance parameters
+        init_params = self._get_init_params()
+        init_params["transforms"] = new_transforms
+
+        # Create new instance
+        new_instance = self.__class__(**init_params)
+
+        # Copy random state from original instance to new instance
+        if hasattr(self, "random_generator") and hasattr(self, "py_random"):
+            new_instance.set_random_state(self.random_generator, self.py_random)
+
+        return new_instance
+
+    def _get_init_params(self) -> dict[str, Any]:
+        """Get parameters needed to recreate this instance.
+
+        Note:
+            Subclasses that add new initialization parameters (other than 'transforms',
+            which is set separately in _create_new_instance) should override this method
+            to include those parameters in the returned dictionary.
+
+        Returns:
+            dict[str, Any]: Dictionary of initialization parameters
+
+        """
+        return {
+            "p": self.p,
+        }
+
+    def _get_effective_seed(self, base_seed: int | None) -> int | None:
+        """Get effective seed considering worker context.
+
+        Args:
+            base_seed (int | None): Base seed value
+
+        Returns:
+            int | None: Effective seed after considering worker context
+
+        """
+        if base_seed is None:
+            return base_seed
+
+        try:
+            import torch
+            import torch.utils.data
+
+            worker_info = torch.utils.data.get_worker_info()
+            if worker_info is not None:
+                # We're in a DataLoader worker process
+                # Use torch.initial_seed() which is unique per worker and changes on respawn
+                torch_seed = torch.initial_seed() % (2**32)
+                return (base_seed + torch_seed) % (2**32)
+        except (ImportError, AttributeError):
+            # PyTorch not available or not in worker context
+            pass
+
+        return base_seed
+
+
+class Compose(BaseCompose, HubMixin):
+    """Compose multiple transforms together and apply them sequentially to input data.
+
+    This class allows you to chain multiple image augmentation transforms and apply them
+    in a specified order. It also handles bounding box and keypoint transformations if
+    the appropriate parameters are provided.
+
+    The Compose class supports dynamic pipeline modification after initialization using
+    mathematical operators. All parameters (bbox_params, keypoint_params, additional_targets,
+    etc.) are preserved when using operators to modify the pipeline.
+
+    Args:
+        transforms (list[BasicTransform | BaseCompose]): A list of transforms to apply.
+        bbox_params (dict[str, Any] | BboxParams | None): Parameters for bounding box transforms.
+            Can be a dict of params or a BboxParams object. Default is None.
+        keypoint_params (dict[str, Any] | KeypointParams | None): Parameters for keypoint transforms.
+            Can be a dict of params or a KeypointParams object. Default is None.
+        additional_targets (dict[str, str] | None): A dictionary mapping additional target names
+            to their types. For example, {'image2': 'image'}. Default is None.
+        p (float): Probability of applying all transforms. Should be in range [0, 1]. Default is 1.0.
+        is_check_shapes (bool): If True, checks consistency of shapes for image/mask/masks on each call.
+            Disable only if you are sure about your data consistency. Default is True.
+        strict (bool): If True, enables strict mode which:
+            1. Validates that all input keys are known/expected
+            2. Validates that no transforms have invalid arguments
+            3. Raises ValueError if any validation fails
+            If False, these validations are skipped. Default is False.
+        mask_interpolation (int | None): Interpolation method for mask transforms. When defined,
+            it overrides the interpolation method specified in individual transforms. Default is None.
+        seed (int | None): Controls reproducibility of random augmentations. Compose uses
+            its own internal random state, completely independent from global random seeds.
+
+            When seed is set (int):
+            - Creates a fixed internal random state
+            - Two Compose instances with the same seed and transforms will produce identical
+              sequences of augmentations
+            - Each call to the same Compose instance still produces random augmentations,
+              but these sequences are reproducible between different Compose instances
+            - Example: transform1 = A.Compose([...], seed=137) and
+                      transform2 = A.Compose([...], seed=137) will produce identical sequences
+
+            When seed is None (default):
+            - Generates a new internal random state on each Compose creation
+            - Different Compose instances will produce different sequences of augmentations
+            - Example: transform = A.Compose([...])  # random results
+
+            Important: Setting random seeds outside of Compose (like np.random.seed() or
+            random.seed()) has no effect on augmentations as Compose uses its own internal
+            random state.
+        save_applied_params (bool): If True, saves the applied parameters of each transform. Default is False.
+            You will need to use the `applied_transforms` key in the output dictionary to access the parameters.
+
+    Examples:
+        >>> # Basic usage:
+        >>> import albumentations as A
+        >>> transform = A.Compose([
+        ...     A.RandomCrop(width=256, height=256),
+        ...     A.HorizontalFlip(p=0.5),
+        ...     A.RandomBrightnessContrast(p=0.2),
+        ... ], seed=137)
+        >>> transformed = transform(image=image)
+
+        >>> # Pipeline modification after initialization:
+        >>> # Create initial pipeline with bbox support
+        >>> base_transform = A.Compose([
+        ...     A.HorizontalFlip(p=0.5),
+        ...     A.RandomCrop(width=512, height=512)
+        ... ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))
+        >>>
+        >>> # Add transforms using operators (bbox_params preserved)
+        >>> extended = base_transform + A.RandomBrightnessContrast(p=0.3)
+        >>> extended = base_transform + [A.Blur(), A.GaussNoise()]
+        >>> extended = A.Resize(height=1024, width=1024) + base_transform
+        >>>
+        >>> # Remove transforms by class
+        >>> pipeline = A.Compose([A.HorizontalFlip(p=0.5), A.VerticalFlip(), A.Rotate()])
+        >>> without_flip = pipeline - A.HorizontalFlip  # Remove by class
+
+    Note:
+        - The class checks the validity of input data and shapes if is_check_args and is_check_shapes are True.
+        - When bbox_params or keypoint_params are provided, it sets up the corresponding processors.
+        - The transform can handle additional targets specified in the additional_targets dictionary.
+        - When strict mode is enabled, it performs additional validation to ensure data and transform
+          configuration correctness.
+        - Pipeline modification operators (+, -, __radd__) preserve all Compose parameters including
+          bbox_params, keypoint_params, additional_targets, and other configuration settings.
+        - All operators return new Compose instances without modifying the original pipeline.
+
+    """
+
+    def __init__(
+        self,
+        transforms: TransformsSeqType,
+        bbox_params: dict[str, Any] | BboxParams | None = None,
+        keypoint_params: dict[str, Any] | KeypointParams | None = None,
+        additional_targets: dict[str, str] | None = None,
+        p: float = 1.0,
+        is_check_shapes: bool = True,
+        strict: bool = False,
+        mask_interpolation: int | None = None,
+        seed: int | None = None,
+        save_applied_params: bool = False,
+    ):
+        # Store the original base seed for worker context recalculation
+        self._base_seed = seed
+
+        # Get effective seed considering worker context
+        effective_seed = self._get_effective_seed(seed)
+
+        super().__init__(
+            transforms=transforms,
+            p=p,
+            mask_interpolation=mask_interpolation,
+            seed=effective_seed,
+            save_applied_params=save_applied_params,
+        )
+
+        if bbox_params:
+            if isinstance(bbox_params, dict):
+                b_params = BboxParams(**bbox_params)
+            elif isinstance(bbox_params, BboxParams):
+                b_params = bbox_params
+            else:
+                msg = "unknown format of bbox_params, please use `dict` or `BboxParams`"
+                raise ValueError(msg)
+            self.processors["bboxes"] = BboxProcessor(b_params)
+
+        if keypoint_params:
+            if isinstance(keypoint_params, dict):
+                k_params = KeypointParams(**keypoint_params)
+            elif isinstance(keypoint_params, KeypointParams):
+                k_params = keypoint_params
+            else:
+                msg = "unknown format of keypoint_params, please use `dict` or `KeypointParams`"
+                raise ValueError(msg)
+            self.processors["keypoints"] = KeypointsProcessor(k_params)
+
+        for proc in self.processors.values():
+            proc.ensure_transforms_valid(self.transforms)
+
+        self.add_targets(additional_targets)
+        if not self.transforms:  # if no transforms -> do nothing, all keys will be available
+            self._available_keys.update(AVAILABLE_KEYS)
+
+        self.is_check_args = True
+        self.strict = strict
+
+        self.is_check_shapes = is_check_shapes
+        self.check_each_transform = tuple(  # processors that checks after each transform
+            proc for proc in self.processors.values() if getattr(proc.params, "check_each_transform", False)
+        )
+        self._set_check_args_for_transforms(self.transforms)
+
+        self._set_processors_for_transforms(self.transforms)
+
+        self.save_applied_params = save_applied_params
+        self._images_was_list = False
+        self._masks_was_list = False
+        self._last_torch_seed: int | None = None
+
+    @property
+    def strict(self) -> bool:
+        """Get the current strict mode setting.
+
+        Returns:
+            bool: True if strict mode is enabled, False otherwise.
+
+        """
+        return self._strict
+
+    @strict.setter
+    def strict(self, value: bool) -> None:
+        # if value and not self._strict:
+        if value:
+            # Only validate when enabling strict mode
+            self._validate_strict()
+        self._strict = value
+
+    def _validate_strict(self) -> None:
+        """Validate that no transforms have invalid arguments when strict mode is enabled."""
+
+        def check_transform(transform: TransformType) -> None:
+            if hasattr(transform, "invalid_args") and transform.invalid_args:
+                message = (
+                    f"Argument(s) '{', '.join(transform.invalid_args)}' "
+                    f"are not valid for transform {transform.__class__.__name__}"
+                )
+                raise ValueError(message)
+            if isinstance(transform, BaseCompose):
+                for t in transform.transforms:
+                    check_transform(t)
+
+        for transform in self.transforms:
+            check_transform(transform)
+
+    def _set_processors_for_transforms(self, transforms: TransformsSeqType) -> None:
+        for transform in transforms:
+            if isinstance(transform, BasicTransform):
+                if hasattr(transform, "set_processors"):
+                    transform.set_processors(self.processors)
+            elif isinstance(transform, BaseCompose):
+                self._set_processors_for_transforms(transform.transforms)
+
+    def _set_check_args_for_transforms(self, transforms: TransformsSeqType) -> None:
+        for transform in transforms:
+            if isinstance(transform, BaseCompose):
+                self._set_check_args_for_transforms(transform.transforms)
+                transform.check_each_transform = self.check_each_transform
+                transform.processors = self.processors
+            if isinstance(transform, Compose):
+                transform.disable_check_args_private()
+
+    def disable_check_args_private(self) -> None:
+        """Disable argument checking for transforms.
+
+        This method disables strict mode and argument checking for all transforms in the composition.
+        """
+        self.is_check_args = False
+        self.strict = False
+        self.main_compose = False
+
+    def __call__(self, *args: Any, force_apply: bool = False, **data: Any) -> dict[str, Any]:
+        """Apply transformations to data with automatic worker seed synchronization.
+
+        Args:
+            *args (Any): Positional arguments are not supported.
+            force_apply (bool): Whether to apply transforms regardless of probability. Default: False.
+            **data (Any): Dict with data to transform.
+
+        Returns:
+            dict[str, Any]: Dictionary with transformed data.
+
+        Raises:
+            KeyError: If positional arguments are provided.
+
+        """
+        # Check and sync worker seed if needed
+        self._check_worker_seed()
+
+        if args:
+            msg = "You have to pass data to augmentations as named arguments, for example: aug(image=image)"
+            raise KeyError(msg)
+
+        # Initialize applied_transforms only in top-level Compose if requested
+        if self.save_applied_params and self.main_compose:
+            data["applied_transforms"] = []
+
+        need_to_run = force_apply or self.py_random.random() < self.p
+        if not need_to_run:
+            return data
+
+        self.preprocess(data)
+
+        for t in self.transforms:
+            data = t(**data)
+            self._track_transform_params(t, data)
+            data = self.check_data_post_transform(data)
+
+        return self.postprocess(data)
+
+    def _check_worker_seed(self) -> None:
+        """Check and update random seed if in worker context."""
+        if not hasattr(self, "_base_seed") or self._base_seed is None:
+            return
+
+        # Check if we're in a worker and need to update the seed
+        try:
+            import torch
+            import torch.utils.data
+
+            worker_info = torch.utils.data.get_worker_info()
+            if worker_info is not None:
+                # Get the current torch initial seed
+                current_torch_seed = torch.initial_seed()
+
+                # Check if we've already synchronized for this seed
+                if hasattr(self, "_last_torch_seed") and self._last_torch_seed == current_torch_seed:
+                    return
+
+                # Update the seed and mark as synchronized
+                self._last_torch_seed = current_torch_seed
+                effective_seed = self._get_effective_seed(self._base_seed)
+
+                # Update our own random state
+                self.random_generator = np.random.default_rng(effective_seed)
+                self.py_random = random.Random(effective_seed)
+
+                # Propagate to all transforms
+                for transform in self.transforms:
+                    if hasattr(transform, "set_random_state"):
+                        transform.set_random_state(self.random_generator, self.py_random)
+                    elif hasattr(transform, "set_random_seed"):
+                        # For transforms that don't have set_random_state, use set_random_seed
+                        transform.set_random_seed(effective_seed)
+        except (ImportError, AttributeError):
+            pass
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Set state from unpickling and handle worker seed."""
+        self.__dict__.update(state)
+        # If we have a base seed, recalculate effective seed in worker context
+        if hasattr(self, "_base_seed") and self._base_seed is not None:
+            # Reset _last_torch_seed to ensure worker-seed sync runs after unpickling
+            self._last_torch_seed = None
+            # Recalculate effective seed in worker context
+            self.set_random_seed(self._base_seed)
+        elif hasattr(self, "seed") and self.seed is not None:
+            # For backward compatibility, if no base seed but seed exists
+            self._base_seed = self.seed
+            self._last_torch_seed = None
+            self.set_random_seed(self.seed)
+
+    def set_random_seed(self, seed: int | None) -> None:
+        """Override to use worker-aware seed functionality.
+
+        Args:
+            seed (int | None): Random seed to use
+
+        """
+        # Store the original base seed
+        self._base_seed = seed
+        self.seed = seed
+
+        # Get effective seed considering worker context
+        effective_seed = self._get_effective_seed(seed)
+
+        # Initialize random generators with effective seed
+        self.random_generator = np.random.default_rng(effective_seed)
+        self.py_random = random.Random(effective_seed)
+
+        # Propagate to all transforms
+        for transform in self.transforms:
+            if hasattr(transform, "set_random_state"):
+                transform.set_random_state(self.random_generator, self.py_random)
+            elif hasattr(transform, "set_random_seed"):
+                # For transforms that don't have set_random_state, use set_random_seed
+                transform.set_random_seed(effective_seed)
+
+    def preprocess(self, data: Any) -> None:
+        """Preprocess input data before applying transforms."""
+        # Always validate shapes if is_check_shapes is True, regardless of strict mode
+        if self.is_check_shapes:
+            shapes = []  # For H,W checks
+            volume_shapes = []  # For D,H,W checks
+
+            for data_name, data_value in data.items():
+                internal_name = self._additional_targets.get(data_name, data_name)
+
+                # Skip empty data
+                if data_value is None:
+                    continue
+
+                shape = self._get_data_shape(data_name, internal_name, data_value)
+                if shape is not None:
+                    if internal_name in CHECKED_VOLUME | CHECKED_MASK3D:
+                        shapes.append(shape[1:3])  # H,W from (D,H,W)
+                        volume_shapes.append(shape[:3])  # D,H,W
+                    elif internal_name in {"volumes", "masks3d"}:
+                        shapes.append(shape[2:4])  # H,W from (N,D,H,W)
+                        volume_shapes.append(shape[1:4])  # D,H,W from (N,D,H,W)
+                    else:
+                        shapes.append(shape[:2])  # H,W
+
+            self._check_shape_consistency(shapes, volume_shapes)
+
+        # Do strict validation only if enabled
+        if self.strict:
+            self._validate_data(data)
+
+        self._preprocess_processors(data)
+        self._preprocess_arrays(data)
+
+    def _validate_data(self, data: dict[str, Any]) -> None:
+        """Validate input data keys and arguments."""
+        if not self.strict:
+            return
+
+        for data_name in data:
+            if not self._is_valid_key(data_name):
+                raise ValueError(f"Key {data_name} is not in available keys.")
+
+        if self.is_check_args:
+            self._check_args(**data)
+
+    def _is_valid_key(self, key: str) -> bool:
+        """Check if the key is valid for processing."""
+        return key in self._available_keys or key in MASK_KEYS or key in IMAGE_KEYS or key == "applied_transforms"
+
+    def _preprocess_processors(self, data: dict[str, Any]) -> None:
+        """Run preprocessors if this is the main compose."""
+        if not self.main_compose:
+            return
+
+        for processor in self.processors.values():
+            processor.ensure_data_valid(data)
+        for processor in self.processors.values():
+            processor.preprocess(data)
+
+    def _preprocess_arrays(self, data: dict[str, Any]) -> None:
+        """Convert lists to numpy arrays for images and masks, and ensure contiguity."""
+        self._preprocess_images(data)
+        self._preprocess_masks(data)
+
+    def _preprocess_images(self, data: dict[str, Any]) -> None:
+        """Convert image lists to numpy arrays."""
+        if "images" not in data:
+            return
+
+        if isinstance(data["images"], (list, tuple)):
+            self._images_was_list = True
+            # Skip stacking for empty lists
+            if not data["images"]:
+                return
+            data["images"] = np.stack(data["images"])
+        else:
+            self._images_was_list = False
+
+    def _preprocess_masks(self, data: dict[str, Any]) -> None:
+        """Convert mask lists to numpy arrays."""
+        if "masks" not in data:
+            return
+
+        if isinstance(data["masks"], (list, tuple)):
+            self._masks_was_list = True
+            # Skip stacking for empty lists
+            if not data["masks"]:
+                return
+            data["masks"] = np.stack(data["masks"])
+        else:
+            self._masks_was_list = False
+
+    def postprocess(self, data: dict[str, Any]) -> dict[str, Any]:
+        """Apply post-processing to data after all transforms have been applied.
+
+        Args:
+            data (dict[str, Any]): Data after transformation.
+
+        Returns:
+            dict[str, Any]: Post-processed data.
+
+        """
+        if self.main_compose:
+            for p in self.processors.values():
+                p.postprocess(data)
+
+            # Convert back to list if original input was a list
+            if "images" in data and self._images_was_list:
+                data["images"] = list(data["images"])
+
+            if "masks" in data and self._masks_was_list:
+                data["masks"] = list(data["masks"])
+
+        return data
+
+    def to_dict_private(self) -> dict[str, Any]:
+        """Convert the composition to a dictionary for serialization.
+
+        Returns:
+            dict[str, Any]: Dictionary representation of the composition.
+
+        """
+        dictionary = super().to_dict_private()
+        bbox_processor = self.processors.get("bboxes")
+        keypoints_processor = self.processors.get("keypoints")
+        dictionary.update(
+            {
+                "bbox_params": bbox_processor.params.to_dict_private() if bbox_processor else None,
+                "keypoint_params": (keypoints_processor.params.to_dict_private() if keypoints_processor else None),
+                "additional_targets": self.additional_targets,
+                "is_check_shapes": self.is_check_shapes,
+                "seed": getattr(self, "_base_seed", None),
+            },
+        )
+        return dictionary
+
+    def get_dict_with_id(self) -> dict[str, Any]:
+        """Get a dictionary representation with object IDs for replay mode.
+
+        Returns:
+            dict[str, Any]: Dictionary with composition data and object IDs.
+
+        """
+        dictionary = super().get_dict_with_id()
+        bbox_processor = self.processors.get("bboxes")
+        keypoints_processor = self.processors.get("keypoints")
+        dictionary.update(
+            {
+                "bbox_params": bbox_processor.params.to_dict_private() if bbox_processor else None,
+                "keypoint_params": (keypoints_processor.params.to_dict_private() if keypoints_processor else None),
+                "additional_targets": self.additional_targets,
+                "params": None,
+                "is_check_shapes": self.is_check_shapes,
+            },
+        )
+        return dictionary
+
+    @staticmethod
+    def _check_single_data(data_name: str, data: Any) -> tuple[int, int]:
+        if not isinstance(data, np.ndarray):
+            raise TypeError(f"{data_name} must be numpy array type")
+        return data.shape[:2]
+
+    @staticmethod
+    def _check_masks_data(data_name: str, data: Any) -> tuple[int, int] | None:
+        """Check masks data format and return shape.
+
+        Args:
+            data_name (str): Name of the data field being checked
+            data (Any): Input data in one of these formats:
+                - List of numpy arrays, each of shape (H, W) or (H, W, C)
+                - Numpy array of shape (N, H, W) or (N, H, W, C)
+                - Empty list for cases where no masks are present
+
+        Returns:
+            tuple[int, int] | None: (height, width) of the first mask, or None if masks list is empty
+        Raises:
+            TypeError: If data format is invalid
+
+        """
+        if isinstance(data, np.ndarray):
+            if data.ndim not in [3, 4]:  # (N,H,W) or (N,H,W,C)
+                raise TypeError(f"{data_name} as numpy array must be 3D or 4D")
+            return data.shape[1:3]  # Return (H,W)
+
+        if isinstance(data, (list, tuple)):
+            if not data:
+                # Allow empty list/tuple of masks
+                return None
+            if not all(isinstance(m, np.ndarray) for m in data):
+                raise TypeError(f"All elements in {data_name} must be numpy arrays")
+            if any(m.ndim not in {2, 3} for m in data):
+                raise TypeError(f"All masks in {data_name} must be 2D or 3D numpy arrays")
+            return data[0].shape[:2]
+
+        raise TypeError(f"{data_name} must be either a numpy array or a sequence of numpy arrays")
+
+    @staticmethod
+    def _check_multi_data(data_name: str, data: Any) -> tuple[int, int]:
+        """Check multi-image data format and return shape.
+
+        Args:
+            data_name (str): Name of the data field being checked
+            data (Any): Input data in one of these formats:
+                - List-like of numpy arrays
+                - Numpy array of shape (N, H, W, C) or (N, H, W)
+
+        Returns:
+            tuple[int, int]: (height, width) of the first image
+        Raises:
+            TypeError: If data format is invalid
+
+        """
+        if isinstance(data, np.ndarray):
+            if data.ndim not in {3, 4}:  # (N,H,W) or (N,H,W,C)
+                raise TypeError(f"{data_name} as numpy array must be 3D or 4D")
+            return data.shape[1:3]  # Return (H,W)
+
+        if not isinstance(data, Sequence) or not isinstance(data[0], np.ndarray):
+            raise TypeError(f"{data_name} must be either a numpy array or a list of numpy arrays")
+        return data[0].shape[:2]
+
+    @staticmethod
+    def _check_bbox_keypoint_params(internal_data_name: str, processors: dict[str, Any]) -> None:
+        if internal_data_name in CHECK_BBOX_PARAM and processors.get("bboxes") is None:
+            raise ValueError("bbox_params must be specified for bbox transformations")
+        if internal_data_name in CHECK_KEYPOINTS_PARAM and processors.get("keypoints") is None:
+            raise ValueError("keypoints_params must be specified for keypoint transformations")
+
+    @staticmethod
+    def _check_shapes(shapes: list[tuple[int, ...]], is_check_shapes: bool) -> None:
+        if is_check_shapes and shapes and shapes.count(shapes[0]) != len(shapes):
+            raise ValueError(
+                "Height and Width of image, mask or masks should be equal. You can disable shapes check "
+                "by setting a parameter is_check_shapes=False of Compose class (do it only if you are sure "
+                "about your data consistency).",
+            )
+
+    def _check_args(self, **kwargs: Any) -> None:
+        shapes = []  # For H,W checks
+        volume_shapes = []  # For D,H,W checks
+
+        for data_name, data in kwargs.items():
+            internal_name = self._additional_targets.get(data_name, data_name)
+
+            # For CHECKED_SINGLE, we must validate even if None
+            if internal_name in CHECKED_SINGLE:
+                if not isinstance(data, np.ndarray):
+                    raise TypeError(f"{data_name} must be numpy array type")
+                shapes.append(data.shape[:2])
+                continue
+
+            # Skip empty data or non-array/list inputs for other types
+            if data is None:
+                continue
+            if not isinstance(data, (np.ndarray, list)):
+                continue
+
+            self._check_bbox_keypoint_params(internal_name, self.processors)
+
+            shape = self._get_data_shape(data_name, internal_name, data)
+            if shape is None:
+                continue
+
+            # Handle different shape types
+            if internal_name in CHECKED_VOLUME | CHECKED_MASK3D:
+                shapes.append(shape[1:3])  # H,W from (D,H,W)
+                volume_shapes.append(shape[:3])  # D,H,W
+            elif internal_name in {"volumes", "masks3d"}:
+                shapes.append(shape[2:4])  # H,W from (N,D,H,W)
+                volume_shapes.append(shape[1:4])  # D,H,W from (N,D,H,W)
+            else:
+                shapes.append(shape[:2])  # H,W
+
+        self._check_shape_consistency(shapes, volume_shapes)
+
+    def _get_data_shape(self, data_name: str, internal_name: str, data: Any) -> tuple[int, ...] | None:
+        """Get shape of data based on its type."""
+        # Handle single images and masks
+        if internal_name in CHECKED_SINGLE:
+            return self._get_single_data_shape(data_name, data)
+
+        # Handle volumes
+        if internal_name in CHECKED_VOLUME:
+            return self._check_volume_data(data_name, data)
+
+        # Handle 3D masks
+        if internal_name in CHECKED_MASK3D:
+            return self._check_mask3d_data(data_name, data)
+
+        # Handle multi-item data (masks, images, volumes)
+        if internal_name in CHECKED_MULTI:
+            return self._get_multi_data_shape(data_name, internal_name, data)
+
+        return None
+
+    def _get_single_data_shape(self, data_name: str, data: np.ndarray) -> tuple[int, ...]:
+        """Get shape of single image or mask."""
+        if not isinstance(data, np.ndarray):
+            raise TypeError(f"{data_name} must be numpy array type")
+        return data.shape
+
+    def _get_multi_data_shape(self, data_name: str, internal_name: str, data: Any) -> tuple[int, ...] | None:
+        """Get shape of multi-item data (masks, images, volumes)."""
+        if internal_name == "masks":
+            shape = self._check_masks_data(data_name, data)
+            # Skip empty masks lists when returning shape
+            return None if shape is None else shape
+
+        if internal_name in {"volumes", "masks3d"}:  # Group these together
+            if not isinstance(data, np.ndarray):
+                raise TypeError(f"{data_name} must be numpy array type")
+            if data.ndim not in {4, 5}:  # (N,D,H,W) or (N,D,H,W,C)
+                raise TypeError(f"{data_name} must be 4D or 5D array")
+            return data.shape  # Return full shape
+
+        return self._check_multi_data(data_name, data)
+
+    def _check_shape_consistency(self, shapes: list[tuple[int, ...]], volume_shapes: list[tuple[int, ...]]) -> None:
+        """Check consistency of shapes."""
+        # Check H,W consistency
+        self._check_shapes(shapes, self.is_check_shapes)
+
+        # Check D,H,W consistency for volumes and 3D masks
+        if self.is_check_shapes and volume_shapes and volume_shapes.count(volume_shapes[0]) != len(volume_shapes):
+            raise ValueError(
+                "Depth, Height and Width of volume, mask3d, volumes and masks3d should be equal. "
+                "You can disable shapes check by setting is_check_shapes=False.",
+            )
+
+    @staticmethod
+    def _check_volume_data(data_name: str, data: np.ndarray) -> tuple[int, int, int]:
+        if data.ndim not in {3, 4}:  # (D,H,W) or (D,H,W,C)
+            raise TypeError(f"{data_name} must be 3D or 4D array")
+        return data.shape[:3]  # Return (D,H,W)
+
+    @staticmethod
+    def _check_volumes_data(data_name: str, data: np.ndarray) -> tuple[int, int, int]:
+        if data.ndim not in {4, 5}:  # (N,D,H,W) or (N,D,H,W,C)
+            raise TypeError(f"{data_name} must be 4D or 5D array")
+        return data.shape[1:4]  # Return (D,H,W)
+
+    @staticmethod
+    def _check_mask3d_data(data_name: str, data: np.ndarray) -> tuple[int, int, int]:
+        """Check single volumetric mask data format and return shape."""
+        if data.ndim not in {3, 4}:  # (D,H,W) or (D,H,W,C)
+            raise TypeError(f"{data_name} must be 3D or 4D array")
+        return data.shape[:3]  # Return (D,H,W)
+
+    @staticmethod
+    def _check_masks3d_data(data_name: str, data: np.ndarray) -> tuple[int, int, int]:
+        """Check multiple volumetric masks data format and return shape."""
+        if data.ndim not in [4, 5]:  # (N,D,H,W) or (N,D,H,W,C)
+            raise TypeError(f"{data_name} must be 4D or 5D array")
+        return data.shape[1:4]  # Return (D,H,W)
+
+    def _get_init_params(self) -> dict[str, Any]:
+        """Get parameters needed to recreate this Compose instance.
+
+        Returns:
+            dict[str, Any]: Dictionary of initialization parameters
+
+        """
+        bbox_processor = self.processors.get("bboxes")
+        keypoints_processor = self.processors.get("keypoints")
+
+        return {
+            "bbox_params": bbox_processor.params if bbox_processor else None,
+            "keypoint_params": keypoints_processor.params if keypoints_processor else None,
+            "additional_targets": self.additional_targets,
+            "p": self.p,
+            "is_check_shapes": self.is_check_shapes,
+            "strict": self.strict,
+            "mask_interpolation": getattr(self, "mask_interpolation", None),
+            "seed": getattr(self, "_base_seed", None),
+            "save_applied_params": getattr(self, "save_applied_params", False),
+        }
+
+
+class OneOf(BaseCompose):
+    """Select one of transforms to apply. Selected transform will be called with `force_apply=True`.
+    Transforms probabilities will be normalized to one 1, so in this case transforms probabilities works as weights.
+
+    Args:
+        transforms (list): list of transformations to compose.
+        p (float): probability of applying selected transform. Default: 0.5.
+
+    """
+
+    def __init__(self, transforms: TransformsSeqType, p: float = 0.5):
+        super().__init__(transforms=transforms, p=p)
+        transforms_ps = [t.p for t in self.transforms]
+        s = sum(transforms_ps)
+        self.transforms_ps = [t / s for t in transforms_ps]
+
+    def __call__(self, *args: Any, force_apply: bool = False, **data: Any) -> dict[str, Any]:
+        """Apply the OneOf composition to the input data.
+
+        Args:
+            *args (Any): Positional arguments are not supported.
+            force_apply (bool): Whether to apply transforms regardless of probability. Default: False.
+            **data (Any): Dict with data to transform.
+
+        Returns:
+            dict[str, Any]: Dictionary with transformed data.
+
+        Raises:
+            KeyError: If positional arguments are provided.
+
+        """
+        if self.replay_mode:
+            for t in self.transforms:
+                data = t(**data)
+            return data
+
+        if self.transforms_ps and (force_apply or self.py_random.random() < self.p):
+            idx: int = self.random_generator.choice(len(self.transforms), p=self.transforms_ps)
+            t = self.transforms[idx]
+            data = t(force_apply=True, **data)
+            self._track_transform_params(t, data)
+        return data
+
+
+class SomeOf(BaseCompose):
+    """Selects exactly `n` transforms from the given list and applies them.
+
+    The selection of which `n` transforms to apply is done **uniformly at random**
+    from the provided list. Each transform in the list has an equal chance of being selected.
+
+    Once the `n` transforms are selected, each one is applied **based on its
+    individual probability** `p`.
+
+    Args:
+        transforms (list[BasicTransform | BaseCompose]): A list of transforms to choose from.
+        n (int): The exact number of transforms to select and potentially apply.
+                 If `replace=False` and `n` is greater than the number of available transforms,
+                 `n` will be capped at the number of transforms.
+        replace (bool): Whether to sample transforms with replacement. If True, the same
+                        transform can be selected multiple times (up to `n` times).
+                        Default is False.
+        p (float): The probability that this `SomeOf` composition will be applied.
+                   If applied, it will select `n` transforms and attempt to apply them.
+                   Default is 1.0.
+
+    Note:
+        - The overall probability `p` of the `SomeOf` block determines if *any* selection
+          and application occurs.
+        - The individual probability `p` of each transform inside the list determines if
+          that specific transform runs *if it is selected*.
+        - If `replace` is True, the same transform might be selected multiple times, and
+          its individual probability `p` will be checked each time it's encountered.
+        - When using pipeline modification operators (+, -, __radd__), the `n` parameter
+          is preserved while the pool of available transforms changes:
+          - `SomeOf([A, B], n=2) + C` → `SomeOf([A, B, C], n=2)` (selects 2 from 3 transforms)
+          - This allows for dynamic adjustment of the transform pool without changing selection count.
+
+    Examples:
+        >>> import albumentations as A
+        >>> transform = A.SomeOf([
+        ...     A.HorizontalFlip(p=0.5),  # 50% chance to apply if selected
+        ...     A.VerticalFlip(p=0.8),    # 80% chance to apply if selected
+        ...     A.RandomRotate90(p=1.0), # 100% chance to apply if selected
+        ... ], n=2, replace=False, p=1.0) # Always select 2 transforms uniformly
+
+        # In each call, 2 transforms out of 3 are chosen uniformly.
+        # For example, if HFlip and VFlip are chosen:
+        # - HFlip runs if random() < 0.5
+        # - VFlip runs if random() < 0.8
+        # If VFlip and Rotate90 are chosen:
+        # - VFlip runs if random() < 0.8
+        # - Rotate90 runs if random() < 1.0 (always)
+
+        >>> # Pipeline modification example:
+        >>> # Add more transforms to the pool while keeping n=2
+        >>> extended = transform + [A.Blur(p=1.0), A.RandomBrightnessContrast(p=0.7)]
+        >>> # Now selects 2 transforms from 5 available transforms uniformly
+
+    """
+
+    def __init__(self, transforms: TransformsSeqType, n: int = 1, replace: bool = False, p: float = 1):
+        super().__init__(transforms, p)
+        self.n = n
+        if not replace and n > len(self.transforms):
+            self.n = len(self.transforms)
+            warnings.warn(
+                f"`n` is greater than number of transforms. `n` will be set to {self.n}.",
+                UserWarning,
+                stacklevel=2,
+            )
+        self.replace = replace
+
+    def __call__(self, *arg: Any, force_apply: bool = False, **data: Any) -> dict[str, Any]:
+        """Apply n randomly selected transforms from the list of transforms.
+
+        Args:
+            *arg (Any): Positional arguments are not supported.
+            force_apply (bool): Whether to apply transforms regardless of probability. Default: False.
+            **data (Any): Dict with data to transform.
+
+        Returns:
+            dict[str, Any]: Dictionary with transformed data.
+
+        """
+        if self.replay_mode:
+            for t in self.transforms:
+                data = t(**data)
+                data = self.check_data_post_transform(data)
+            return data
+
+        if self.py_random.random() < self.p:  # Check overall SomeOf probability
+            # Get indices uniformly
+            indices_to_consider = self._get_idx()
+            for i in indices_to_consider:
+                t = self.transforms[i]
+                # Apply the transform respecting its own probability `t.p`
+                data = t(**data)
+                self._track_transform_params(t, data)
+                data = self.check_data_post_transform(data)
+        return data
+
+    def _get_idx(self) -> np.ndarray[np.int_]:
+        # Use uniform probability for selection, ignore individual p values here
+        idx = self.random_generator.choice(
+            len(self.transforms),
+            size=self.n,
+            replace=self.replace,
+        )
+        idx.sort()
+        return idx
+
+    def to_dict_private(self) -> dict[str, Any]:
+        """Convert the SomeOf composition to a dictionary for serialization.
+
+        Returns:
+            dict[str, Any]: Dictionary representation of the composition.
+
+        """
+        dictionary = super().to_dict_private()
+        dictionary.update({"n": self.n, "replace": self.replace})
+        return dictionary
+
+    def _get_init_params(self) -> dict[str, Any]:
+        """Get parameters needed to recreate this SomeOf instance.
+
+        Returns:
+            dict[str, Any]: Dictionary of initialization parameters
+
+        """
+        base_params = super()._get_init_params()
+        base_params.update(
+            {
+                "n": self.n,
+                "replace": self.replace,
+            },
+        )
+        return base_params
+
+
+class RandomOrder(SomeOf):
+    """Apply a random subset of transforms from the given list in a random order.
+
+    Selects exactly `n` transforms uniformly at random from the list, and then applies
+    the selected transforms in a random order. Each selected transform is applied
+    based on its individual probability `p`.
+
+    Attributes:
+        transforms (TransformsSeqType): A list of transformations to choose from.
+        n (int): The number of transforms to apply. If `n` is greater than the number of available transforms
+                 and `replace` is False, `n` will be set to the number of available transforms.
+        replace (bool): Whether to sample transforms with replacement. If True, the same transform can be
+                        selected multiple times. Default is False.
+        p (float): Probability of applying the selected transforms. Should be in the range [0, 1]. Default is 1.0.
+
+    Examples:
+        >>> import albumentations as A
+        >>> transform = A.RandomOrder([
+        ...     A.HorizontalFlip(p=0.5),
+        ...     A.VerticalFlip(p=1.0),
+        ...     A.RandomBrightnessContrast(p=0.8),
+        ... ], n=2, replace=False, p=1.0)
+        >>> # This will uniformly select 2 transforms and apply them in a random order,
+        >>> # respecting their individual probabilities (0.5, 1.0, 0.8).
+
+    Note:
+        - Inherits from SomeOf, but overrides `_get_idx` to ensure random order without sorting.
+        - Selection is uniform; application depends on individual transform probabilities.
+
+    """
+
+    def __init__(self, transforms: TransformsSeqType, n: int = 1, replace: bool = False, p: float = 1):
+        # Initialize using SomeOf's logic (which now does uniform selection setup)
+        super().__init__(transforms=transforms, n=n, replace=replace, p=p)
+
+    def _get_idx(self) -> np.ndarray[np.int_]:
+        # Perform uniform random selection without replacement, like SomeOf
+        # Crucially, DO NOT sort the indices here to maintain random order.
+        return self.random_generator.choice(
+            len(self.transforms),
+            size=self.n,
+            replace=self.replace,
+        )
+
+
+class OneOrOther(BaseCompose):
+    """Select one or another transform to apply. Selected transform will be called with `force_apply=True`."""
+
+    def __init__(
+        self,
+        first: TransformType | None = None,
+        second: TransformType | None = None,
+        transforms: TransformsSeqType | None = None,
+        p: float = 0.5,
+    ):
+        if transforms is None:
+            if first is None or second is None:
+                msg = "You must set both first and second or set transforms argument."
+                raise ValueError(msg)
+            transforms = [first, second]
+        super().__init__(transforms=transforms, p=p)
+        if len(self.transforms) != NUM_ONEOF_TRANSFORMS:
+            warnings.warn("Length of transforms is not equal to 2.", stacklevel=2)
+
+    def __call__(self, *args: Any, force_apply: bool = False, **data: Any) -> dict[str, Any]:
+        """Apply one or another transform to the input data.
+
+        Args:
+            *args (Any): Positional arguments are not supported.
+            force_apply (bool): Whether to apply transforms regardless of probability. Default: False.
+            **data (Any): Dict with data to transform.
+
+        Returns:
+            dict[str, Any]: Dictionary with transformed data.
+
+        """
+        if self.replay_mode:
+            for t in self.transforms:
+                data = t(**data)
+                self._track_transform_params(t, data)
+            return data
+
+        if self.py_random.random() < self.p:
+            return self.transforms[0](force_apply=True, **data)
+
+        return self.transforms[-1](force_apply=True, **data)
+
+
+class SelectiveChannelTransform(BaseCompose):
+    """A transformation class to apply specified transforms to selected channels of an image.
+
+    This class extends BaseCompose to allow selective application of transformations to
+    specified image channels. It extracts the selected channels, applies the transformations,
+    and then reinserts the transformed channels back into their original positions in the image.
+
+    Args:
+        transforms (TransformsSeqType):
+            A sequence of transformations (from Albumentations) to be applied to the specified channels.
+        channels (Sequence[int]):
+            A sequence of integers specifying the indices of the channels to which the transforms should be applied.
+        p (float): Probability that the transform will be applied; the default is 1.0 (always apply).
+
+    Returns:
+        dict[str, Any]: The transformed data dictionary, which includes the transformed 'image' key.
+
+    Note:
+        - When using pipeline modification operators (+, -, __radd__), the `channels` parameter
+          is preserved in the resulting SelectiveChannelTransform instance.
+        - Only the transform list is modified while maintaining the same channel selection behavior.
+
+    """
+
+    def __init__(
+        self,
+        transforms: TransformsSeqType,
+        channels: Sequence[int] = (0, 1, 2),
+        p: float = 1.0,
+    ) -> None:
+        super().__init__(transforms=transforms, p=p)
+        self.channels = channels
+
+    def __call__(self, *args: Any, force_apply: bool = False, **data: Any) -> dict[str, Any]:
+        """Apply transforms to specific channels of the image.
+
+        Args:
+            *args (Any): Positional arguments are not supported.
+            force_apply (bool): Whether to apply transforms regardless of probability. Default: False.
+            **data (Any): Dict with data to transform.
+
+        Returns:
+            dict[str, Any]: Dictionary with transformed data.
+
+        """
+        if force_apply or self.py_random.random() < self.p:
+            image = data["image"]
+
+            selected_channels = image[:, :, self.channels]
+            sub_image = np.ascontiguousarray(selected_channels)
+
+            for t in self.transforms:
+                sub_data = {"image": sub_image}
+                sub_image = t(**sub_data)["image"]
+                self._track_transform_params(t, sub_data)
+
+            transformed_channels = cv2.split(sub_image)
+            output_img = image.copy()
+
+            for idx, channel in zip(self.channels, transformed_channels):
+                output_img[:, :, idx] = channel
+
+            data["image"] = np.ascontiguousarray(output_img)
+
+        return data
+
+    def _get_init_params(self) -> dict[str, Any]:
+        """Get parameters needed to recreate this SelectiveChannelTransform instance.
+
+        Returns:
+            dict[str, Any]: Dictionary of initialization parameters
+
+        """
+        base_params = super()._get_init_params()
+        base_params.update(
+            {
+                "channels": self.channels,
+            },
+        )
+        return base_params
+
+
+class ReplayCompose(Compose):
+    """Composition class that enables transform replay functionality.
+
+    This class extends the Compose class with the ability to record and replay
+    transformations. This is useful for applying the same sequence of random
+    transformations to different data.
+
+    Args:
+        transforms (TransformsSeqType): List of transformations to compose.
+        bbox_params (dict[str, Any] | BboxParams | None): Parameters for bounding box transforms.
+        keypoint_params (dict[str, Any] | KeypointParams | None): Parameters for keypoint transforms.
+        additional_targets (dict[str, str] | None): Dictionary of additional targets.
+        p (float): Probability of applying the compose.
+        is_check_shapes (bool): Whether to check shapes of different targets.
+        save_key (str): Key for storing the applied transformations.
+
+    """
+
+    def __init__(
+        self,
+        transforms: TransformsSeqType,
+        bbox_params: dict[str, Any] | BboxParams | None = None,
+        keypoint_params: dict[str, Any] | KeypointParams | None = None,
+        additional_targets: dict[str, str] | None = None,
+        p: float = 1.0,
+        is_check_shapes: bool = True,
+        save_key: str = "replay",
+    ):
+        super().__init__(transforms, bbox_params, keypoint_params, additional_targets, p, is_check_shapes)
+        self.set_deterministic(True, save_key=save_key)
+        self.save_key = save_key
+        self._available_keys.add(save_key)
+
+    def __call__(self, *args: Any, force_apply: bool = False, **kwargs: Any) -> dict[str, Any]:
+        """Apply transforms and record parameters for future replay.
+
+        Args:
+            *args (Any): Positional arguments are not supported.
+            force_apply (bool): Whether to apply transforms regardless of probability. Default: False.
+            **kwargs (Any): Dict with data to transform.
+
+        Returns:
+            dict[str, Any]: Dictionary with transformed data and replay information.
+
+        """
+        kwargs[self.save_key] = defaultdict(dict)
+        result = super().__call__(force_apply=force_apply, **kwargs)
+        serialized = self.get_dict_with_id()
+        self.fill_with_params(serialized, result[self.save_key])
+        self.fill_applied(serialized)
+        result[self.save_key] = serialized
+        return result
+
+    @staticmethod
+    def replay(saved_augmentations: dict[str, Any], **kwargs: Any) -> dict[str, Any]:
+        """Replay previously saved augmentations.
+
+        Args:
+            saved_augmentations (dict[str, Any]): Previously saved augmentation parameters.
+            **kwargs (Any): Dict with data to transform.
+
+        Returns:
+            dict[str, Any]: Dictionary with transformed data using saved parameters.
+
+        """
+        augs = ReplayCompose._restore_for_replay(saved_augmentations)
+        return augs(force_apply=True, **kwargs)
+
+    @staticmethod
+    def _restore_for_replay(
+        transform_dict: dict[str, Any],
+        lambda_transforms: dict[str, Any] | None = None,
+    ) -> TransformType:
+        """Args:
+        transform_dict (dict[str, Any]): A dictionary that contains transform data.
+        lambda_transforms (dict): A dictionary that contains lambda transforms, that
+            is instances of the Lambda class.
+        This dictionary is required when you are restoring a pipeline that contains lambda transforms.
+        Keys in that dictionary should be named same as `name` arguments in respective lambda transforms
+        from a serialized pipeline.
+
+        """
+        applied = transform_dict["applied"]
+        params = transform_dict["params"]
+        lmbd = instantiate_nonserializable(transform_dict, lambda_transforms)
+        if lmbd:
+            transform = lmbd
+        else:
+            name = transform_dict["__class_fullname__"]
+            args = {k: v for k, v in transform_dict.items() if k not in ["__class_fullname__", "applied", "params"]}
+            cls = SERIALIZABLE_REGISTRY[name]
+            if "transforms" in args:
+                args["transforms"] = [
+                    ReplayCompose._restore_for_replay(t, lambda_transforms=lambda_transforms)
+                    for t in args["transforms"]
+                ]
+            transform = cls(**args)
+
+        transform = cast("BasicTransform", transform)
+        if isinstance(transform, BasicTransform):
+            transform.params = params
+        transform.replay_mode = True
+        transform.applied_in_replay = applied
+        return transform
+
+    def fill_with_params(self, serialized: dict[str, Any], all_params: Any) -> None:
+        """Fill serialized transform data with parameters for replay.
+
+        Args:
+            serialized (dict[str, Any]): Serialized transform data.
+            all_params (Any): Parameters to fill in.
+
+        """
+        params = all_params.get(serialized.get("id"))
+        serialized["params"] = params
+        del serialized["id"]
+        for transform in serialized.get("transforms", []):
+            self.fill_with_params(transform, all_params)
+
+    def fill_applied(self, serialized: dict[str, Any]) -> bool:
+        """Set 'applied' flag for transforms based on parameters.
+
+        Args:
+            serialized (dict[str, Any]): Serialized transform data.
+
+        Returns:
+            bool: True if any transform was applied, False otherwise.
+
+        """
+        if "transforms" in serialized:
+            applied = [self.fill_applied(t) for t in serialized["transforms"]]
+            serialized["applied"] = any(applied)
+        else:
+            serialized["applied"] = serialized.get("params") is not None
+        return serialized["applied"]
+
+    def to_dict_private(self) -> dict[str, Any]:
+        """Convert the ReplayCompose to a dictionary for serialization.
+
+        Returns:
+            dict[str, Any]: Dictionary representation of the composition.
+
+        """
+        dictionary = super().to_dict_private()
+        dictionary.update({"save_key": self.save_key})
+        return dictionary
+
+    def _get_init_params(self) -> dict[str, Any]:
+        """Get parameters needed to recreate this ReplayCompose instance.
+
+        Returns:
+            dict[str, Any]: Dictionary of initialization parameters
+
+        """
+        base_params = super()._get_init_params()
+        base_params.update(
+            {
+                "save_key": self.save_key,
+            },
+        )
+        return base_params
+
+
+class Sequential(BaseCompose):
+    """Sequentially applies all transforms to targets.
+
+    Note:
+        This transform is not intended to be a replacement for `Compose`. Instead, it should be used inside `Compose`
+        the same way `OneOf` or `OneOrOther` are used. For instance, you can combine `OneOf` with `Sequential` to
+        create an augmentation pipeline that contains multiple sequences of augmentations and applies one randomly
+        chose sequence to input data (see the `Example` section for an example definition of such pipeline).
+
+    Examples:
+        >>> import albumentations as A
+        >>> transform = A.Compose([
+        >>>    A.OneOf([
+        >>>        A.Sequential([
+        >>>            A.HorizontalFlip(p=0.5),
+        >>>            A.ShiftScaleRotate(p=0.5),
+        >>>        ]),
+        >>>        A.Sequential([
+        >>>            A.VerticalFlip(p=0.5),
+        >>>            A.RandomBrightnessContrast(p=0.5),
+        >>>        ]),
+        >>>    ], p=1)
+        >>> ])
+
+    """
+
+    def __init__(self, transforms: TransformsSeqType, p: float = 0.5):
+        super().__init__(transforms=transforms, p=p)
+
+    def __call__(self, *args: Any, force_apply: bool = False, **data: Any) -> dict[str, Any]:
+        """Apply all transforms in sequential order.
+
+        Args:
+            *args (Any): Positional arguments are not supported.
+            force_apply (bool): Whether to apply transforms regardless of probability. Default: False.
+            **data (Any): Dict with data to transform.
+
+        Returns:
+            dict[str, Any]: Dictionary with transformed data.
+
+        """
+        if self.replay_mode or force_apply or self.py_random.random() < self.p:
+            for t in self.transforms:
+                data = t(**data)
+                self._track_transform_params(t, data)
+                data = self.check_data_post_transform(data)
+        return data