ultralytics 8.1.29__py3-none-any.whl → 8.3.63__py3-none-any.whl

ultralytics/data/augment.py

@@ -1,107 +1,388 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import math
 import random
 from copy import deepcopy
+from typing import Tuple, Union
 
 import cv2
 import numpy as np
 import torch
-import torchvision.transforms as T
+from PIL import Image
 
+from ultralytics.data.utils import polygons2masks, polygons2masks_overlap
 from ultralytics.utils import LOGGER, colorstr
 from ultralytics.utils.checks import check_version
 from ultralytics.utils.instance import Instances
 from ultralytics.utils.metrics import bbox_ioa
 from ultralytics.utils.ops import segment2box, xyxyxyxy2xywhr
 from ultralytics.utils.torch_utils import TORCHVISION_0_10, TORCHVISION_0_11, TORCHVISION_0_13
-from .utils import polygons2masks, polygons2masks_overlap
 
 DEFAULT_MEAN = (0.0, 0.0, 0.0)
 DEFAULT_STD = (1.0, 1.0, 1.0)
-DEFAULT_CROP_FTACTION = 1.0
+DEFAULT_CROP_FRACTION = 1.0
 
 
-# TODO: we might need a BaseTransform to make all these augments be compatible with both classification and semantic
 class BaseTransform:
     """
-    Base class for image transformations.
+    Base class for image transformations in the Ultralytics library.
 
-    This is a generic transformation class that can be extended for specific image processing needs.
-    The class is designed to be compatible with both classification and semantic segmentation tasks.
+    This class serves as a foundation for implementing various image processing operations, designed to be
+    compatible with both classification and semantic segmentation tasks.
 
     Methods:
-        __init__: Initializes the BaseTransform object.
-        apply_image: Applies image transformation to labels.
+        apply_image: Applies image transformations to labels.
         apply_instances: Applies transformations to object instances in labels.
         apply_semantic: Applies semantic segmentation to an image.
         __call__: Applies all label transformations to an image, instances, and semantic masks.
+
+    Examples:
+        >>> transform = BaseTransform()
+        >>> labels = {"image": np.array(...), "instances": [...], "semantic": np.array(...)}
+        >>> transformed_labels = transform(labels)
     """
 
     def __init__(self) -> None:
-        """Initializes the BaseTransform object."""
+        """
+        Initializes the BaseTransform object.
+
+        This constructor sets up the base transformation object, which can be extended for specific image
+        processing tasks. It is designed to be compatible with both classification and semantic segmentation.
+
+        Examples:
+            >>> transform = BaseTransform()
+        """
         pass
 
     def apply_image(self, labels):
-        """Applies image transformations to labels."""
+        """
+        Applies image transformations to labels.
+
+        This method is intended to be overridden by subclasses to implement specific image transformation
+        logic. In its base form, it returns the input labels unchanged.
+
+        Args:
+            labels (Any): The input labels to be transformed. The exact type and structure of labels may
+                vary depending on the specific implementation.
+
+        Returns:
+            (Any): The transformed labels. In the base implementation, this is identical to the input.
+
+        Examples:
+            >>> transform = BaseTransform()
+            >>> original_labels = [1, 2, 3]
+            >>> transformed_labels = transform.apply_image(original_labels)
+            >>> print(transformed_labels)
+            [1, 2, 3]
+        """
         pass
 
     def apply_instances(self, labels):
-        """Applies transformations to object instances in labels."""
+        """
+        Applies transformations to object instances in labels.
+
+        This method is responsible for applying various transformations to object instances within the given
+        labels. It is designed to be overridden by subclasses to implement specific instance transformation
+        logic.
+
+        Args:
+            labels (Dict): A dictionary containing label information, including object instances.
+
+        Returns:
+            (Dict): The modified labels dictionary with transformed object instances.
+
+        Examples:
+            >>> transform = BaseTransform()
+            >>> labels = {"instances": Instances(xyxy=torch.rand(5, 4), cls=torch.randint(0, 80, (5,)))}
+            >>> transformed_labels = transform.apply_instances(labels)
+        """
         pass
 
     def apply_semantic(self, labels):
-        """Applies semantic segmentation to an image."""
+        """
+        Applies semantic segmentation transformations to an image.
+
+        This method is intended to be overridden by subclasses to implement specific semantic segmentation
+        transformations. In its base form, it does not perform any operations.
+
+        Args:
+            labels (Any): The input labels or semantic segmentation mask to be transformed.
+
+        Returns:
+            (Any): The transformed semantic segmentation mask or labels.
+
+        Examples:
+            >>> transform = BaseTransform()
+            >>> semantic_mask = np.zeros((100, 100), dtype=np.uint8)
+            >>> transformed_mask = transform.apply_semantic(semantic_mask)
+        """
         pass
 
     def __call__(self, labels):
-        """Applies all label transformations to an image, instances, and semantic masks."""
+        """
+        Applies all label transformations to an image, instances, and semantic masks.
+
+        This method orchestrates the application of various transformations defined in the BaseTransform class
+        to the input labels. It sequentially calls the apply_image and apply_instances methods to process the
+        image and object instances, respectively.
+
+        Args:
+            labels (Dict): A dictionary containing image data and annotations. Expected keys include 'img' for
+                the image data, and 'instances' for object instances.
+
+        Returns:
+            (Dict): The input labels dictionary with transformed image and instances.
+
+        Examples:
+            >>> transform = BaseTransform()
+            >>> labels = {"img": np.random.rand(640, 640, 3), "instances": []}
+            >>> transformed_labels = transform(labels)
+        """
         self.apply_image(labels)
         self.apply_instances(labels)
         self.apply_semantic(labels)
 
 
 class Compose:
-    """Class for composing multiple image transformations."""
+    """
+    A class for composing multiple image transformations.
+
+    Attributes:
+        transforms (List[Callable]): A list of transformation functions to be applied sequentially.
+
+    Methods:
+        __call__: Applies a series of transformations to input data.
+        append: Appends a new transform to the existing list of transforms.
+        insert: Inserts a new transform at a specified index in the list of transforms.
+        __getitem__: Retrieves a specific transform or a set of transforms using indexing.
+        __setitem__: Sets a specific transform or a set of transforms using indexing.
+        tolist: Converts the list of transforms to a standard Python list.
+
+    Examples:
+        >>> transforms = [RandomFlip(), RandomPerspective(30)]
+        >>> compose = Compose(transforms)
+        >>> transformed_data = compose(data)
+        >>> compose.append(CenterCrop((224, 224)))
+        >>> compose.insert(0, RandomFlip())
+    """
 
     def __init__(self, transforms):
-        """Initializes the Compose object with a list of transforms."""
-        self.transforms = transforms
+        """
+        Initializes the Compose object with a list of transforms.
+
+        Args:
+            transforms (List[Callable]): A list of callable transform objects to be applied sequentially.
+
+        Examples:
+            >>> from ultralytics.data.augment import Compose, RandomHSV, RandomFlip
+            >>> transforms = [RandomHSV(), RandomFlip()]
+            >>> compose = Compose(transforms)
+        """
+        self.transforms = transforms if isinstance(transforms, list) else [transforms]
 
     def __call__(self, data):
-        """Applies a series of transformations to input data."""
+        """
+        Applies a series of transformations to input data. This method sequentially applies each transformation in the
+        Compose object's list of transforms to the input data.
+
+        Args:
+            data (Any): The input data to be transformed. This can be of any type, depending on the
+                transformations in the list.
+
+        Returns:
+            (Any): The transformed data after applying all transformations in sequence.
+
+        Examples:
+            >>> transforms = [Transform1(), Transform2(), Transform3()]
+            >>> compose = Compose(transforms)
+            >>> transformed_data = compose(input_data)
+        """
         for t in self.transforms:
             data = t(data)
         return data
 
     def append(self, transform):
-        """Appends a new transform to the existing list of transforms."""
+        """
+        Appends a new transform to the existing list of transforms.
+
+        Args:
+            transform (BaseTransform): The transformation to be added to the composition.
+
+        Examples:
+            >>> compose = Compose([RandomFlip(), RandomPerspective()])
+            >>> compose.append(RandomHSV())
+        """
         self.transforms.append(transform)
 
+    def insert(self, index, transform):
+        """
+        Inserts a new transform at a specified index in the existing list of transforms.
+
+        Args:
+            index (int): The index at which to insert the new transform.
+            transform (BaseTransform): The transform object to be inserted.
+
+        Examples:
+            >>> compose = Compose([Transform1(), Transform2()])
+            >>> compose.insert(1, Transform3())
+            >>> len(compose.transforms)
+            3
+        """
+        self.transforms.insert(index, transform)
+
+    def __getitem__(self, index: Union[list, int]) -> "Compose":
+        """
+        Retrieves a specific transform or a set of transforms using indexing.
+
+        Args:
+            index (int | List[int]): Index or list of indices of the transforms to retrieve.
+
+        Returns:
+            (Compose): A new Compose object containing the selected transform(s).
+
+        Raises:
+            AssertionError: If the index is not of type int or list.
+
+        Examples:
+            >>> transforms = [RandomFlip(), RandomPerspective(10), RandomHSV(0.5, 0.5, 0.5)]
+            >>> compose = Compose(transforms)
+            >>> single_transform = compose[1]  # Returns a Compose object with only RandomPerspective
+            >>> multiple_transforms = compose[0:2]  # Returns a Compose object with RandomFlip and RandomPerspective
+        """
+        assert isinstance(index, (int, list)), f"The indices should be either list or int type but got {type(index)}"
+        index = [index] if isinstance(index, int) else index
+        return Compose([self.transforms[i] for i in index])
+
+    def __setitem__(self, index: Union[list, int], value: Union[list, int]) -> None:
+        """
+        Sets one or more transforms in the composition using indexing.
+
+        Args:
+            index (int | List[int]): Index or list of indices to set transforms at.
+            value (Any | List[Any]): Transform or list of transforms to set at the specified index(es).
+
+        Raises:
+            AssertionError: If index type is invalid, value type doesn't match index type, or index is out of range.
+
+        Examples:
+            >>> compose = Compose([Transform1(), Transform2(), Transform3()])
+            >>> compose[1] = NewTransform()  # Replace second transform
+            >>> compose[0:2] = [NewTransform1(), NewTransform2()]  # Replace first two transforms
+        """
+        assert isinstance(index, (int, list)), f"The indices should be either list or int type but got {type(index)}"
+        if isinstance(index, list):
+            assert isinstance(value, list), (
+                f"The indices should be the same type as values, but got {type(index)} and {type(value)}"
+            )
+        if isinstance(index, int):
+            index, value = [index], [value]
+        for i, v in zip(index, value):
+            assert i < len(self.transforms), f"list index {i} out of range {len(self.transforms)}."
+            self.transforms[i] = v
+
     def tolist(self):
-        """Converts the list of transforms to a standard Python list."""
+        """
+        Converts the list of transforms to a standard Python list.
+
+        Returns:
+            (List): A list containing all the transform objects in the Compose instance.
+
+        Examples:
+            >>> transforms = [RandomFlip(), RandomPerspective(10), CenterCrop()]
+            >>> compose = Compose(transforms)
+            >>> transform_list = compose.tolist()
+            >>> print(len(transform_list))
+            3
+        """
         return self.transforms
 
     def __repr__(self):
-        """Returns a string representation of the object."""
+        """
+        Returns a string representation of the Compose object.
+
+        Returns:
+            (str): A string representation of the Compose object, including the list of transforms.
+
+        Examples:
+            >>> transforms = [RandomFlip(), RandomPerspective(degrees=10, translate=0.1, scale=0.1)]
+            >>> compose = Compose(transforms)
+            >>> print(compose)
+            Compose([
+                RandomFlip(),
+                RandomPerspective(degrees=10, translate=0.1, scale=0.1)
+            ])
+        """
         return f"{self.__class__.__name__}({', '.join([f'{t}' for t in self.transforms])})"
 
 
 class BaseMixTransform:
     """
-    Class for base mix (MixUp/Mosaic) transformations.
+    Base class for mix transformations like MixUp and Mosaic.
+
+    This class provides a foundation for implementing mix transformations on datasets. It handles the
+    probability-based application of transforms and manages the mixing of multiple images and labels.
 
-    This implementation is from mmyolo.
+    Attributes:
+        dataset (Any): The dataset object containing images and labels.
+        pre_transform (Callable | None): Optional transform to apply before mixing.
+        p (float): Probability of applying the mix transformation.
+
+    Methods:
+        __call__: Applies the mix transformation to the input labels.
+        _mix_transform: Abstract method to be implemented by subclasses for specific mix operations.
+        get_indexes: Abstract method to get indexes of images to be mixed.
+        _update_label_text: Updates label text for mixed images.
+
+    Examples:
+        >>> class CustomMixTransform(BaseMixTransform):
+        ...     def _mix_transform(self, labels):
+        ...         # Implement custom mix logic here
+        ...         return labels
+        ...
+        ...     def get_indexes(self):
+        ...         return [random.randint(0, len(self.dataset) - 1) for _ in range(3)]
+        >>> dataset = YourDataset()
+        >>> transform = CustomMixTransform(dataset, p=0.5)
+        >>> mixed_labels = transform(original_labels)
     """
 
     def __init__(self, dataset, pre_transform=None, p=0.0) -> None:
-        """Initializes the BaseMixTransform object with dataset, pre_transform, and probability."""
+        """
+        Initializes the BaseMixTransform object for mix transformations like MixUp and Mosaic.
+
+        This class serves as a base for implementing mix transformations in image processing pipelines.
+
+        Args:
+            dataset (Any): The dataset object containing images and labels for mixing.
+            pre_transform (Callable | None): Optional transform to apply before mixing.
+            p (float): Probability of applying the mix transformation. Should be in the range [0.0, 1.0].
+
+        Examples:
+            >>> dataset = YOLODataset("path/to/data")
+            >>> pre_transform = Compose([RandomFlip(), RandomPerspective()])
+            >>> mix_transform = BaseMixTransform(dataset, pre_transform, p=0.5)
+        """
         self.dataset = dataset
         self.pre_transform = pre_transform
         self.p = p
 
     def __call__(self, labels):
-        """Applies pre-processing transforms and mixup/mosaic transforms to labels data."""
+        """
+        Applies pre-processing transforms and mixup/mosaic transforms to labels data.
+
+        This method determines whether to apply the mix transform based on a probability factor. If applied, it
+        selects additional images, applies pre-transforms if specified, and then performs the mix transform.
+
+        Args:
+            labels (Dict): A dictionary containing label data for an image.
+
+        Returns:
+            (Dict): The transformed labels dictionary, which may include mixed data from other images.
+
+        Examples:
+            >>> transform = BaseMixTransform(dataset, pre_transform=None, p=0.5)
+            >>> result = transform({"image": img, "bboxes": boxes, "cls": classes})
+        """
         if random.uniform(0, 1) > self.p:
             return labels
 
@@ -118,53 +399,197 @@ class BaseMixTransform:
                 mix_labels[i] = self.pre_transform(data)
         labels["mix_labels"] = mix_labels
 
+        # Update cls and texts
+        labels = self._update_label_text(labels)
         # Mosaic or MixUp
         labels = self._mix_transform(labels)
         labels.pop("mix_labels", None)
         return labels
 
     def _mix_transform(self, labels):
-        """Applies MixUp or Mosaic augmentation to the label dictionary."""
+        """
+        Applies MixUp or Mosaic augmentation to the label dictionary.
+
+        This method should be implemented by subclasses to perform specific mix transformations like MixUp or
+        Mosaic. It modifies the input label dictionary in-place with the augmented data.
+
+        Args:
+            labels (Dict): A dictionary containing image and label data. Expected to have a 'mix_labels' key
+                with a list of additional image and label data for mixing.
+
+        Returns:
+            (Dict): The modified labels dictionary with augmented data after applying the mix transform.
+
+        Examples:
+            >>> transform = BaseMixTransform(dataset)
+            >>> labels = {"image": img, "bboxes": boxes, "mix_labels": [{"image": img2, "bboxes": boxes2}]}
+            >>> augmented_labels = transform._mix_transform(labels)
+        """
         raise NotImplementedError
 
     def get_indexes(self):
-        """Gets a list of shuffled indexes for mosaic augmentation."""
+        """
+        Gets a list of shuffled indexes for mosaic augmentation.
+
+        Returns:
+            (List[int]): A list of shuffled indexes from the dataset.
+
+        Examples:
+            >>> transform = BaseMixTransform(dataset)
+            >>> indexes = transform.get_indexes()
+            >>> print(indexes)  # [3, 18, 7, 2]
+        """
         raise NotImplementedError
 
+    @staticmethod
+    def _update_label_text(labels):
+        """
+        Updates label text and class IDs for mixed labels in image augmentation.
+
+        This method processes the 'texts' and 'cls' fields of the input labels dictionary and any mixed labels,
+        creating a unified set of text labels and updating class IDs accordingly.
+
+        Args:
+            labels (Dict): A dictionary containing label information, including 'texts' and 'cls' fields,
+                and optionally a 'mix_labels' field with additional label dictionaries.
+
+        Returns:
+            (Dict): The updated labels dictionary with unified text labels and updated class IDs.
+
+        Examples:
+            >>> labels = {
+            ...     "texts": [["cat"], ["dog"]],
+            ...     "cls": torch.tensor([[0], [1]]),
+            ...     "mix_labels": [{"texts": [["bird"], ["fish"]], "cls": torch.tensor([[0], [1]])}],
+            ... }
+            >>> updated_labels = self._update_label_text(labels)
+            >>> print(updated_labels["texts"])
+            [['cat'], ['dog'], ['bird'], ['fish']]
+            >>> print(updated_labels["cls"])
+            tensor([[0],
+                    [1]])
+            >>> print(updated_labels["mix_labels"][0]["cls"])
+            tensor([[2],
+                    [3]])
+        """
+        if "texts" not in labels:
+            return labels
+
+        mix_texts = sum([labels["texts"]] + [x["texts"] for x in labels["mix_labels"]], [])
+        mix_texts = list({tuple(x) for x in mix_texts})
+        text2id = {text: i for i, text in enumerate(mix_texts)}
+
+        for label in [labels] + labels["mix_labels"]:
+            for i, cls in enumerate(label["cls"].squeeze(-1).tolist()):
+                text = label["texts"][int(cls)]
+                label["cls"][i] = text2id[tuple(text)]
+            label["texts"] = mix_texts
+        return labels
+
 
 class Mosaic(BaseMixTransform):
     """
-    Mosaic augmentation.
+    Mosaic augmentation for image datasets.
 
     This class performs mosaic augmentation by combining multiple (4 or 9) images into a single mosaic image.
     The augmentation is applied to a dataset with a given probability.
 
     Attributes:
         dataset: The dataset on which the mosaic augmentation is applied.
-        imgsz (int, optional): Image size (height and width) after mosaic pipeline of a single image. Default to 640.
-        p (float, optional): Probability of applying the mosaic augmentation. Must be in the range 0-1. Default to 1.0.
-        n (int, optional): The grid size, either 4 (for 2x2) or 9 (for 3x3).
+        imgsz (int): Image size (height and width) after mosaic pipeline of a single image.
+        p (float): Probability of applying the mosaic augmentation. Must be in the range 0-1.
+        n (int): The grid size, either 4 (for 2x2) or 9 (for 3x3).
+        border (Tuple[int, int]): Border size for width and height.
+
+    Methods:
+        get_indexes: Returns a list of random indexes from the dataset.
+        _mix_transform: Applies mixup transformation to the input image and labels.
+        _mosaic3: Creates a 1x3 image mosaic.
+        _mosaic4: Creates a 2x2 image mosaic.
+        _mosaic9: Creates a 3x3 image mosaic.
+        _update_labels: Updates labels with padding.
+        _cat_labels: Concatenates labels and clips mosaic border instances.
+
+    Examples:
+        >>> from ultralytics.data.augment import Mosaic
+        >>> dataset = YourDataset(...)  # Your image dataset
+        >>> mosaic_aug = Mosaic(dataset, imgsz=640, p=0.5, n=4)
+        >>> augmented_labels = mosaic_aug(original_labels)
     """
 
     def __init__(self, dataset, imgsz=640, p=1.0, n=4):
-        """Initializes the object with a dataset, image size, probability, and border."""
+        """
+        Initializes the Mosaic augmentation object.
+
+        This class performs mosaic augmentation by combining multiple (4 or 9) images into a single mosaic image.
+        The augmentation is applied to a dataset with a given probability.
+
+        Args:
+            dataset (Any): The dataset on which the mosaic augmentation is applied.
+            imgsz (int): Image size (height and width) after mosaic pipeline of a single image.
+            p (float): Probability of applying the mosaic augmentation. Must be in the range 0-1.
+            n (int): The grid size, either 4 (for 2x2) or 9 (for 3x3).
+
+        Examples:
+            >>> from ultralytics.data.augment import Mosaic
+            >>> dataset = YourDataset(...)
+            >>> mosaic_aug = Mosaic(dataset, imgsz=640, p=0.5, n=4)
+        """
         assert 0 <= p <= 1.0, f"The probability should be in range [0, 1], but got {p}."
-        assert n in (4, 9), "grid must be equal to 4 or 9."
+        assert n in {4, 9}, "grid must be equal to 4 or 9."
         super().__init__(dataset=dataset, p=p)
-        self.dataset = dataset
         self.imgsz = imgsz
         self.border = (-imgsz // 2, -imgsz // 2)  # width, height
         self.n = n
 
     def get_indexes(self, buffer=True):
-        """Return a list of random indexes from the dataset."""
+        """
+        Returns a list of random indexes from the dataset for mosaic augmentation.
+
+        This method selects random image indexes either from a buffer or from the entire dataset, depending on
+        the 'buffer' parameter. It is used to choose images for creating mosaic augmentations.
+
+        Args:
+            buffer (bool): If True, selects images from the dataset buffer. If False, selects from the entire
+                dataset.
+
+        Returns:
+            (List[int]): A list of random image indexes. The length of the list is n-1, where n is the number
+                of images used in the mosaic (either 3 or 8, depending on whether n is 4 or 9).
+
+        Examples:
+            >>> mosaic = Mosaic(dataset, imgsz=640, p=1.0, n=4)
+            >>> indexes = mosaic.get_indexes()
+            >>> print(len(indexes))  # Output: 3
+        """
         if buffer:  # select images from buffer
             return random.choices(list(self.dataset.buffer), k=self.n - 1)
         else:  # select any images
             return [random.randint(0, len(self.dataset) - 1) for _ in range(self.n - 1)]
 
     def _mix_transform(self, labels):
-        """Apply mixup transformation to the input image and labels."""
+        """
+        Applies mosaic augmentation to the input image and labels.
+
+        This method combines multiple images (3, 4, or 9) into a single mosaic image based on the 'n' attribute.
+        It ensures that rectangular annotations are not present and that there are other images available for
+        mosaic augmentation.
+
+        Args:
+            labels (Dict): A dictionary containing image data and annotations. Expected keys include:
+                - 'rect_shape': Should be None as rect and mosaic are mutually exclusive.
+                - 'mix_labels': A list of dictionaries containing data for other images to be used in the mosaic.
+
+        Returns:
+            (Dict): A dictionary containing the mosaic-augmented image and updated annotations.
+
+        Raises:
+            AssertionError: If 'rect_shape' is not None or if 'mix_labels' is empty.
+
+        Examples:
+            >>> mosaic = Mosaic(dataset, imgsz=640, p=1.0, n=4)
+            >>> augmented_data = mosaic._mix_transform(labels)
+        """
         assert labels.get("rect_shape", None) is None, "rect and mosaic are mutually exclusive."
         assert len(labels.get("mix_labels", [])), "There are no other images for mosaic augment."
         return (
@@ -172,7 +597,32 @@ class Mosaic(BaseMixTransform):
         )  # This code is modified for mosaic3 method.
 
     def _mosaic3(self, labels):
-        """Create a 1x3 image mosaic."""
+        """
+        Creates a 1x3 image mosaic by combining three images.
+
+        This method arranges three images in a horizontal layout, with the main image in the center and two
+        additional images on either side. It's part of the Mosaic augmentation technique used in object detection.
+
+        Args:
+            labels (Dict): A dictionary containing image and label information for the main (center) image.
+                Must include 'img' key with the image array, and 'mix_labels' key with a list of two
+                dictionaries containing information for the side images.
+
+        Returns:
+            (Dict): A dictionary with the mosaic image and updated labels. Keys include:
+                - 'img' (np.ndarray): The mosaic image array with shape (H, W, C).
+                - Other keys from the input labels, updated to reflect the new image dimensions.
+
+        Examples:
+            >>> mosaic = Mosaic(dataset, imgsz=640, p=1.0, n=3)
+            >>> labels = {
+            ...     "img": np.random.rand(480, 640, 3),
+            ...     "mix_labels": [{"img": np.random.rand(480, 640, 3)} for _ in range(2)],
+            ... }
+            >>> result = mosaic._mosaic3(labels)
+            >>> print(result["img"].shape)
+            (640, 640, 3)
+        """
         mosaic_labels = []
         s = self.imgsz
         for i in range(3):
@@ -192,7 +642,7 @@ class Mosaic(BaseMixTransform):
                 c = s - w, s + h0 - h, s, s + h0
 
             padw, padh = c[:2]
-            x1, y1, x2, y2 = (max(x, 0) for x in c)  # allocate coords
+            x1, y1, x2, y2 = (max(x, 0) for x in c)  # allocate coordinates
 
             img3[y1:y2, x1:x2] = img[y1 - padh :, x1 - padw :]  # img3[ymin:ymax, xmin:xmax]
             # hp, wp = h, w  # height, width previous for next iteration
@@ -206,7 +656,29 @@ class Mosaic(BaseMixTransform):
         return final_labels
 
     def _mosaic4(self, labels):
-        """Create a 2x2 image mosaic."""
+        """
+        Creates a 2x2 image mosaic from four input images.
+
+        This method combines four images into a single mosaic image by placing them in a 2x2 grid. It also
+        updates the corresponding labels for each image in the mosaic.
+
+        Args:
+            labels (Dict): A dictionary containing image data and labels for the base image (index 0) and three
+                additional images (indices 1-3) in the 'mix_labels' key.
+
+        Returns:
+            (Dict): A dictionary containing the mosaic image and updated labels. The 'img' key contains the mosaic
+                image as a numpy array, and other keys contain the combined and adjusted labels for all four images.
+
+        Examples:
+            >>> mosaic = Mosaic(dataset, imgsz=640, p=1.0, n=4)
+            >>> labels = {
+            ...     "img": np.random.rand(480, 640, 3),
+            ...     "mix_labels": [{"img": np.random.rand(480, 640, 3)} for _ in range(3)],
+            ... }
+            >>> result = mosaic._mosaic4(labels)
+            >>> assert result["img"].shape == (1280, 1280, 3)
+        """
         mosaic_labels = []
         s = self.imgsz
         yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.border)  # mosaic center x, y
@@ -242,7 +714,31 @@ class Mosaic(BaseMixTransform):
         return final_labels
 
     def _mosaic9(self, labels):
-        """Create a 3x3 image mosaic."""
+        """
+        Creates a 3x3 image mosaic from the input image and eight additional images.
+
+        This method combines nine images into a single mosaic image. The input image is placed at the center,
+        and eight additional images from the dataset are placed around it in a 3x3 grid pattern.
+
+        Args:
+            labels (Dict): A dictionary containing the input image and its associated labels. It should have
+                the following keys:
+                - 'img' (numpy.ndarray): The input image.
+                - 'resized_shape' (Tuple[int, int]): The shape of the resized image (height, width).
+                - 'mix_labels' (List[Dict]): A list of dictionaries containing information for the additional
+                  eight images, each with the same structure as the input labels.
+
+        Returns:
+            (Dict): A dictionary containing the mosaic image and updated labels. It includes the following keys:
+                - 'img' (numpy.ndarray): The final mosaic image.
+                - Other keys from the input labels, updated to reflect the new mosaic arrangement.
+
+        Examples:
+            >>> mosaic = Mosaic(dataset, imgsz=640, p=1.0, n=9)
+            >>> input_labels = dataset[0]
+            >>> mosaic_result = mosaic._mosaic9(input_labels)
+            >>> mosaic_image = mosaic_result["img"]
+        """
         mosaic_labels = []
         s = self.imgsz
         hp, wp = -1, -1  # height, width previous
@@ -275,7 +771,7 @@ class Mosaic(BaseMixTransform):
                 c = s - w, s + h0 - hp - h, s, s + h0 - hp
 
             padw, padh = c[:2]
-            x1, y1, x2, y2 = (max(x, 0) for x in c)  # allocate coords
+            x1, y1, x2, y2 = (max(x, 0) for x in c)  # allocate coordinates
 
             # Image
             img9[y1:y2, x1:x2] = img[y1 - padh :, x1 - padw :]  # img9[ymin:ymax, xmin:xmax]
@@ -291,7 +787,25 @@ class Mosaic(BaseMixTransform):
 
     @staticmethod
     def _update_labels(labels, padw, padh):
-        """Update labels."""
+        """
+        Updates label coordinates with padding values.
+
+        This method adjusts the bounding box coordinates of object instances in the labels by adding padding
+        values. It also denormalizes the coordinates if they were previously normalized.
+
+        Args:
+            labels (Dict): A dictionary containing image and instance information.
+            padw (int): Padding width to be added to the x-coordinates.
+            padh (int): Padding height to be added to the y-coordinates.
+
+        Returns:
+            (Dict): Updated labels dictionary with adjusted instance coordinates.
+
+        Examples:
+            >>> labels = {"img": np.zeros((100, 100, 3)), "instances": Instances(...)}
+            >>> padw, padh = 50, 50
+            >>> updated_labels = Mosaic._update_labels(labels, padw, padh)
+        """
         nh, nw = labels["img"].shape[:2]
         labels["instances"].convert_bbox(format="xyxy")
         labels["instances"].denormalize(nw, nh)
@@ -299,7 +813,32 @@ class Mosaic(BaseMixTransform):
         return labels
 
     def _cat_labels(self, mosaic_labels):
-        """Return labels with mosaic border instances clipped."""
+        """
+        Concatenates and processes labels for mosaic augmentation.
+
+        This method combines labels from multiple images used in mosaic augmentation, clips instances to the
+        mosaic border, and removes zero-area boxes.
+
+        Args:
+            mosaic_labels (List[Dict]): A list of label dictionaries for each image in the mosaic.
+
+        Returns:
+            (Dict): A dictionary containing concatenated and processed labels for the mosaic image, including:
+                - im_file (str): File path of the first image in the mosaic.
+                - ori_shape (Tuple[int, int]): Original shape of the first image.
+                - resized_shape (Tuple[int, int]): Shape of the mosaic image (imgsz * 2, imgsz * 2).
+                - cls (np.ndarray): Concatenated class labels.
+                - instances (Instances): Concatenated instance annotations.
+                - mosaic_border (Tuple[int, int]): Mosaic border size.
+                - texts (List[str], optional): Text labels if present in the original labels.
+
+        Examples:
+            >>> mosaic = Mosaic(dataset, imgsz=640)
+            >>> mosaic_labels = [{"cls": np.array([0, 1]), "instances": Instances(...)} for _ in range(4)]
+            >>> result = mosaic._cat_labels(mosaic_labels)
+            >>> print(result.keys())
+            dict_keys(['im_file', 'ori_shape', 'resized_shape', 'cls', 'instances', 'mosaic_border'])
+        """
         if len(mosaic_labels) == 0:
             return {}
         cls = []
@@ -320,22 +859,88 @@ class Mosaic(BaseMixTransform):
         final_labels["instances"].clip(imgsz, imgsz)
         good = final_labels["instances"].remove_zero_area_boxes()
         final_labels["cls"] = final_labels["cls"][good]
+        if "texts" in mosaic_labels[0]:
+            final_labels["texts"] = mosaic_labels[0]["texts"]
         return final_labels
 
 
 class MixUp(BaseMixTransform):
-    """Class for applying MixUp augmentation to the dataset."""
+    """
+    Applies MixUp augmentation to image datasets.
+
+    This class implements the MixUp augmentation technique as described in the paper "mixup: Beyond Empirical Risk
+    Minimization" (https://arxiv.org/abs/1710.09412). MixUp combines two images and their labels using a random weight.
+
+    Attributes:
+        dataset (Any): The dataset to which MixUp augmentation will be applied.
+        pre_transform (Callable | None): Optional transform to apply before MixUp.
+        p (float): Probability of applying MixUp augmentation.
+
+    Methods:
+        get_indexes: Returns a random index from the dataset.
+        _mix_transform: Applies MixUp augmentation to the input labels.
+
+    Examples:
+        >>> from ultralytics.data.augment import MixUp
+        >>> dataset = YourDataset(...)  # Your image dataset
+        >>> mixup = MixUp(dataset, p=0.5)
+        >>> augmented_labels = mixup(original_labels)
+    """
 
     def __init__(self, dataset, pre_transform=None, p=0.0) -> None:
-        """Initializes MixUp object with dataset, pre_transform, and probability of applying MixUp."""
+        """
+        Initializes the MixUp augmentation object.
+
+        MixUp is an image augmentation technique that combines two images by taking a weighted sum of their pixel
+        values and labels. This implementation is designed for use with the Ultralytics YOLO framework.
+
+        Args:
+            dataset (Any): The dataset to which MixUp augmentation will be applied.
+            pre_transform (Callable | None): Optional transform to apply to images before MixUp.
+            p (float): Probability of applying MixUp augmentation to an image. Must be in the range [0, 1].
+
+        Examples:
+            >>> from ultralytics.data.dataset import YOLODataset
+            >>> dataset = YOLODataset("path/to/data.yaml")
+            >>> mixup = MixUp(dataset, pre_transform=None, p=0.5)
+        """
         super().__init__(dataset=dataset, pre_transform=pre_transform, p=p)
 
     def get_indexes(self):
-        """Get a random index from the dataset."""
+        """
+        Get a random index from the dataset.
+
+        This method returns a single random index from the dataset, which is used to select an image for MixUp
+        augmentation.
+
+        Returns:
+            (int): A random integer index within the range of the dataset length.
+
+        Examples:
+            >>> mixup = MixUp(dataset)
+            >>> index = mixup.get_indexes()
+            >>> print(index)
+            42
+        """
         return random.randint(0, len(self.dataset) - 1)
 
     def _mix_transform(self, labels):
-        """Applies MixUp augmentation as per https://arxiv.org/pdf/1710.09412.pdf."""
+        """
+        Applies MixUp augmentation to the input labels.
+
+        This method implements the MixUp augmentation technique as described in the paper
+        "mixup: Beyond Empirical Risk Minimization" (https://arxiv.org/abs/1710.09412).
+
+        Args:
+            labels (Dict): A dictionary containing the original image and label information.
+
+        Returns:
+            (Dict): A dictionary containing the mixed-up image and combined label information.
+
+        Examples:
+            >>> mixer = MixUp(dataset)
+            >>> mixed_labels = mixer._mix_transform(labels)
+        """
         r = np.random.beta(32.0, 32.0)  # mixup ratio, alpha=beta=32.0
         labels2 = labels["mix_labels"][0]
         labels["img"] = (labels["img"] * r + labels2["img"] * (1 - r)).astype(np.uint8)
@@ -346,33 +951,61 @@ class MixUp(BaseMixTransform):
 
 class RandomPerspective:
     """
-    Implements random perspective and affine transformations on images and corresponding bounding boxes, segments, and
-    keypoints. These transformations include rotation, translation, scaling, and shearing. The class also offers the
-    option to apply these transformations conditionally with a specified probability.
+    Implements random perspective and affine transformations on images and corresponding annotations.
+
+    This class applies random rotations, translations, scaling, shearing, and perspective transformations
+    to images and their associated bounding boxes, segments, and keypoints. It can be used as part of an
+    augmentation pipeline for object detection and instance segmentation tasks.
 
     Attributes:
-        degrees (float): Degree range for random rotations.
-        translate (float): Fraction of total width and height for random translation.
-        scale (float): Scaling factor interval, e.g., a scale factor of 0.1 allows a resize between 90%-110%.
-        shear (float): Shear intensity (angle in degrees).
+        degrees (float): Maximum absolute degree range for random rotations.
+        translate (float): Maximum translation as a fraction of the image size.
+        scale (float): Scaling factor range, e.g., scale=0.1 means 0.9-1.1.
+        shear (float): Maximum shear angle in degrees.
         perspective (float): Perspective distortion factor.
-        border (tuple): Tuple specifying mosaic border.
-        pre_transform (callable): A function/transform to apply to the image before starting the random transformation.
+        border (Tuple[int, int]): Mosaic border size as (x, y).
+        pre_transform (Callable | None): Optional transform to apply before the random perspective.
 
     Methods:
-        affine_transform(img, border): Applies a series of affine transformations to the image.
-        apply_bboxes(bboxes, M): Transforms bounding boxes using the calculated affine matrix.
-        apply_segments(segments, M): Transforms segments and generates new bounding boxes.
-        apply_keypoints(keypoints, M): Transforms keypoints.
-        __call__(labels): Main method to apply transformations to both images and their corresponding annotations.
-        box_candidates(box1, box2): Filters out bounding boxes that don't meet certain criteria post-transformation.
+        affine_transform: Applies affine transformations to the input image.
+        apply_bboxes: Transforms bounding boxes using the affine matrix.
+        apply_segments: Transforms segments and generates new bounding boxes.
+        apply_keypoints: Transforms keypoints using the affine matrix.
+        __call__: Applies the random perspective transformation to images and annotations.
+        box_candidates: Filters transformed bounding boxes based on size and aspect ratio.
+
+    Examples:
+        >>> transform = RandomPerspective(degrees=10, translate=0.1, scale=0.1, shear=10)
+        >>> image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
+        >>> labels = {"img": image, "cls": np.array([0, 1]), "instances": Instances(...)}
+        >>> result = transform(labels)
+        >>> transformed_image = result["img"]
+        >>> transformed_instances = result["instances"]
     """
 
     def __init__(
         self, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, border=(0, 0), pre_transform=None
     ):
-        """Initializes RandomPerspective object with transformation parameters."""
+        """
+        Initializes RandomPerspective object with transformation parameters.
 
+        This class implements random perspective and affine transformations on images and corresponding bounding boxes,
+        segments, and keypoints. Transformations include rotation, translation, scaling, and shearing.
+
+        Args:
+            degrees (float): Degree range for random rotations.
+            translate (float): Fraction of total width and height for random translation.
+            scale (float): Scaling factor interval, e.g., a scale factor of 0.5 allows a resize between 50%-150%.
+            shear (float): Shear intensity (angle in degrees).
+            perspective (float): Perspective distortion factor.
+            border (Tuple[int, int]): Tuple specifying mosaic border (top/bottom, left/right).
+            pre_transform (Callable | None): Function/transform to apply to the image before starting the random
+                transformation.
+
+        Examples:
+            >>> transform = RandomPerspective(degrees=10.0, translate=0.1, scale=0.5, shear=5.0)
+            >>> result = transform(labels)  # Apply random perspective to labels
+        """
         self.degrees = degrees
         self.translate = translate
         self.scale = scale
@@ -385,16 +1018,26 @@ class RandomPerspective:
         """
         Applies a sequence of affine transformations centered around the image center.
 
+        This function performs a series of geometric transformations on the input image, including
+        translation, perspective change, rotation, scaling, and shearing. The transformations are
+        applied in a specific order to maintain consistency.
+
         Args:
-            img (ndarray): Input image.
-            border (tuple): Border dimensions.
+            img (np.ndarray): Input image to be transformed.
+            border (Tuple[int, int]): Border dimensions for the transformed image.
 
         Returns:
-            img (ndarray): Transformed image.
-            M (ndarray): Transformation matrix.
-            s (float): Scale factor.
+            (Tuple[np.ndarray, np.ndarray, float]): A tuple containing:
+                - np.ndarray: Transformed image.
+                - np.ndarray: 3x3 transformation matrix.
+                - float: Scale factor applied during the transformation.
+
+        Examples:
+            >>> import numpy as np
+            >>> img = np.random.rand(100, 100, 3)
+            >>> border = (10, 10)
+            >>> transformed_img, matrix, scale = affine_transform(img, border)
         """
-
         # Center
         C = np.eye(3, dtype=np.float32)
 
@@ -436,14 +1079,23 @@ class RandomPerspective:
 
     def apply_bboxes(self, bboxes, M):
         """
-        Apply affine to bboxes only.
+        Apply affine transformation to bounding boxes.
+
+        This function applies an affine transformation to a set of bounding boxes using the provided
+        transformation matrix.
 
         Args:
-            bboxes (ndarray): list of bboxes, xyxy format, with shape (num_bboxes, 4).
-            M (ndarray): affine matrix.
+            bboxes (torch.Tensor): Bounding boxes in xyxy format with shape (N, 4), where N is the number
+                of bounding boxes.
+            M (torch.Tensor): Affine transformation matrix with shape (3, 3).
 
         Returns:
-            new_bboxes (ndarray): bboxes after affine, [num_bboxes, 4].
+            (torch.Tensor): Transformed bounding boxes in xyxy format with shape (N, 4).
+
+        Examples:
+            >>> bboxes = torch.tensor([[10, 10, 20, 20], [30, 30, 40, 40]])
+            >>> M = torch.eye(3)
+            >>> transformed_bboxes = apply_bboxes(bboxes, M)
         """
         n = len(bboxes)
         if n == 0:
@@ -461,15 +1113,25 @@ class RandomPerspective:
 
     def apply_segments(self, segments, M):
         """
-        Apply affine to segments and generate new bboxes from segments.
+        Apply affine transformations to segments and generate new bounding boxes.
+
+        This function applies affine transformations to input segments and generates new bounding boxes based on
+        the transformed segments. It clips the transformed segments to fit within the new bounding boxes.
 
         Args:
-            segments (ndarray): list of segments, [num_samples, 500, 2].
-            M (ndarray): affine matrix.
+            segments (np.ndarray): Input segments with shape (N, M, 2), where N is the number of segments and M is the
+                number of points in each segment.
+            M (np.ndarray): Affine transformation matrix with shape (3, 3).
 
         Returns:
-            new_segments (ndarray): list of segments after affine, [num_samples, 500, 2].
-            new_bboxes (ndarray): bboxes after affine, [N, 4].
+            (Tuple[np.ndarray, np.ndarray]): A tuple containing:
+                - New bounding boxes with shape (N, 4) in xyxy format.
+                - Transformed and clipped segments with shape (N, M, 2).
+
+        Examples:
+            >>> segments = np.random.rand(10, 500, 2)  # 10 segments with 500 points each
+            >>> M = np.eye(3)  # Identity transformation matrix
+            >>> new_bboxes, new_segments = apply_segments(segments, M)
         """
         n, num = segments.shape[:2]
         if n == 0:
@@ -488,14 +1150,25 @@ class RandomPerspective:
 
     def apply_keypoints(self, keypoints, M):
         """
-        Apply affine to keypoints.
+        Applies affine transformation to keypoints.
+
+        This method transforms the input keypoints using the provided affine transformation matrix. It handles
+        perspective rescaling if necessary and updates the visibility of keypoints that fall outside the image
+        boundaries after transformation.
 
         Args:
-            keypoints (ndarray): keypoints, [N, 17, 3].
-            M (ndarray): affine matrix.
+            keypoints (np.ndarray): Array of keypoints with shape (N, 17, 3), where N is the number of instances,
+                17 is the number of keypoints per instance, and 3 represents (x, y, visibility).
+            M (np.ndarray): 3x3 affine transformation matrix.
 
         Returns:
-            new_keypoints (ndarray): keypoints after affine, [N, 17, 3].
+            (np.ndarray): Transformed keypoints array with the same shape as input (N, 17, 3).
+
+        Examples:
+            >>> random_perspective = RandomPerspective()
+            >>> keypoints = np.random.rand(5, 17, 3)  # 5 instances, 17 keypoints each
+            >>> M = np.eye(3)  # Identity transformation
+            >>> transformed_keypoints = random_perspective.apply_keypoints(keypoints, M)
         """
         n, nkpt = keypoints.shape[:2]
         if n == 0:
@@ -511,10 +1184,38 @@ class RandomPerspective:
 
     def __call__(self, labels):
         """
-        Affine images and targets.
+        Applies random perspective and affine transformations to an image and its associated labels.
+
+        This method performs a series of transformations including rotation, translation, scaling, shearing,
+        and perspective distortion on the input image and adjusts the corresponding bounding boxes, segments,
+        and keypoints accordingly.
 
         Args:
-            labels (dict): a dict of `bboxes`, `segments`, `keypoints`.
+            labels (Dict): A dictionary containing image data and annotations.
+                Must include:
+                    'img' (ndarray): The input image.
+                    'cls' (ndarray): Class labels.
+                    'instances' (Instances): Object instances with bounding boxes, segments, and keypoints.
+                May include:
+                    'mosaic_border' (Tuple[int, int]): Border size for mosaic augmentation.
+
+        Returns:
+            (Dict): Transformed labels dictionary containing:
+                - 'img' (np.ndarray): The transformed image.
+                - 'cls' (np.ndarray): Updated class labels.
+                - 'instances' (Instances): Updated object instances.
+                - 'resized_shape' (Tuple[int, int]): New image shape after transformation.
+
+        Examples:
+            >>> transform = RandomPerspective()
+            >>> image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
+            >>> labels = {
+            ...     "img": image,
+            ...     "cls": np.array([0, 1, 2]),
+            ...     "instances": Instances(bboxes=np.array([[10, 10, 50, 50], [100, 100, 150, 150]])),
+            ... }
+            >>> result = transform(labels)
+            >>> assert result["img"].shape[:2] == result["resized_shape"]
         """
         if self.pre_transform and "mosaic_border" not in labels:
             labels = self.pre_transform(labels)
@@ -559,21 +1260,39 @@ class RandomPerspective:
         labels["resized_shape"] = img.shape[:2]
         return labels
 
-    def box_candidates(self, box1, box2, wh_thr=2, ar_thr=100, area_thr=0.1, eps=1e-16):
+    @staticmethod
+    def box_candidates(box1, box2, wh_thr=2, ar_thr=100, area_thr=0.1, eps=1e-16):
         """
-        Compute box candidates based on a set of thresholds. This method compares the characteristics of the boxes
-        before and after augmentation to decide whether a box is a candidate for further processing.
+        Compute candidate boxes for further processing based on size and aspect ratio criteria.
+
+        This method compares boxes before and after augmentation to determine if they meet specified
+        thresholds for width, height, aspect ratio, and area. It's used to filter out boxes that have
+        been overly distorted or reduced by the augmentation process.
 
         Args:
-            box1 (numpy.ndarray): The 4,n bounding box before augmentation, represented as [x1, y1, x2, y2].
-            box2 (numpy.ndarray): The 4,n bounding box after augmentation, represented as [x1, y1, x2, y2].
-            wh_thr (float, optional): The width and height threshold in pixels. Default is 2.
-            ar_thr (float, optional): The aspect ratio threshold. Default is 100.
-            area_thr (float, optional): The area ratio threshold. Default is 0.1.
-            eps (float, optional): A small epsilon value to prevent division by zero. Default is 1e-16.
+            box1 (numpy.ndarray): Original boxes before augmentation, shape (4, N) where n is the
+                number of boxes. Format is [x1, y1, x2, y2] in absolute coordinates.
+            box2 (numpy.ndarray): Augmented boxes after transformation, shape (4, N). Format is
+                [x1, y1, x2, y2] in absolute coordinates.
+            wh_thr (float): Width and height threshold in pixels. Boxes smaller than this in either
+                dimension are rejected.
+            ar_thr (float): Aspect ratio threshold. Boxes with an aspect ratio greater than this
+                value are rejected.
+            area_thr (float): Area ratio threshold. Boxes with an area ratio (new/old) less than
+                this value are rejected.
+            eps (float): Small epsilon value to prevent division by zero.
 
         Returns:
-            (numpy.ndarray): A boolean array indicating which boxes are candidates based on the given thresholds.
+            (numpy.ndarray): Boolean array of shape (n) indicating which boxes are candidates.
+                True values correspond to boxes that meet all criteria.
+
+        Examples:
+            >>> random_perspective = RandomPerspective()
+            >>> box1 = np.array([[0, 0, 100, 100], [0, 0, 50, 50]]).T
+            >>> box2 = np.array([[10, 10, 90, 90], [5, 5, 45, 45]]).T
+            >>> candidates = random_perspective.box_candidates(box1, box2)
+            >>> print(candidates)
+            [True True]
         """
         w1, h1 = box1[2] - box1[0], box1[3] - box1[1]
         w2, h2 = box2[2] - box2[0], box2[3] - box2[1]
@@ -583,20 +1302,42 @@ class RandomPerspective:
 
 class RandomHSV:
     """
-    This class is responsible for performing random adjustments to the Hue, Saturation, and Value (HSV) channels of an
-    image.
+    Randomly adjusts the Hue, Saturation, and Value (HSV) channels of an image.
+
+    This class applies random HSV augmentation to images within predefined limits set by hgain, sgain, and vgain.
+
+    Attributes:
+        hgain (float): Maximum variation for hue. Range is typically [0, 1].
+        sgain (float): Maximum variation for saturation. Range is typically [0, 1].
+        vgain (float): Maximum variation for value. Range is typically [0, 1].
 
-    The adjustments are random but within limits set by hgain, sgain, and vgain.
+    Methods:
+        __call__: Applies random HSV augmentation to an image.
+
+    Examples:
+        >>> import numpy as np
+        >>> from ultralytics.data.augment import RandomHSV
+        >>> augmenter = RandomHSV(hgain=0.5, sgain=0.5, vgain=0.5)
+        >>> image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
+        >>> labels = {"img": image}
+        >>> augmenter(labels)
+        >>> augmented_image = augmented_labels["img"]
     """
 
     def __init__(self, hgain=0.5, sgain=0.5, vgain=0.5) -> None:
         """
-        Initialize RandomHSV class with gains for each HSV channel.
+        Initializes the RandomHSV object for random HSV (Hue, Saturation, Value) augmentation.
+
+        This class applies random adjustments to the HSV channels of an image within specified limits.
 
         Args:
-            hgain (float, optional): Maximum variation for hue. Default is 0.5.
-            sgain (float, optional): Maximum variation for saturation. Default is 0.5.
-            vgain (float, optional): Maximum variation for value. Default is 0.5.
+            hgain (float): Maximum variation for hue. Should be in the range [0, 1].
+            sgain (float): Maximum variation for saturation. Should be in the range [0, 1].
+            vgain (float): Maximum variation for value. Should be in the range [0, 1].
+
+        Examples:
+            >>> hsv_aug = RandomHSV(hgain=0.5, sgain=0.5, vgain=0.5)
+            >>> hsv_aug(image)
         """
         self.hgain = hgain
         self.sgain = sgain
@@ -604,9 +1345,24 @@ class RandomHSV:
 
     def __call__(self, labels):
         """
-        Applies random HSV augmentation to an image within the predefined limits.
+        Applies random HSV augmentation to an image within predefined limits.
 
-        The modified image replaces the original image in the input 'labels' dict.
+        This method modifies the input image by randomly adjusting its Hue, Saturation, and Value (HSV) channels.
+        The adjustments are made within the limits set by hgain, sgain, and vgain during initialization.
+
+        Args:
+            labels (Dict): A dictionary containing image data and metadata. Must include an 'img' key with
+                the image as a numpy array.
+
+        Returns:
+            (None): The function modifies the input 'labels' dictionary in-place, updating the 'img' key
+                with the HSV-augmented image.
+
+        Examples:
+            >>> hsv_augmenter = RandomHSV(hgain=0.5, sgain=0.5, vgain=0.5)
+            >>> labels = {"img": np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)}
+            >>> hsv_augmenter(labels)
+            >>> augmented_img = labels["img"]
         """
         img = labels["img"]
         if self.hgain or self.sgain or self.vgain:
@@ -628,21 +1384,45 @@ class RandomFlip:
     """
     Applies a random horizontal or vertical flip to an image with a given probability.
 
-    Also updates any instances (bounding boxes, keypoints, etc.) accordingly.
+    This class performs random image flipping and updates corresponding instance annotations such as
+    bounding boxes and keypoints.
+
+    Attributes:
+        p (float): Probability of applying the flip. Must be between 0 and 1.
+        direction (str): Direction of flip, either 'horizontal' or 'vertical'.
+        flip_idx (array-like): Index mapping for flipping keypoints, if applicable.
+
+    Methods:
+        __call__: Applies the random flip transformation to an image and its annotations.
+
+    Examples:
+        >>> transform = RandomFlip(p=0.5, direction="horizontal")
+        >>> result = transform({"img": image, "instances": instances})
+        >>> flipped_image = result["img"]
+        >>> flipped_instances = result["instances"]
     """
 
     def __init__(self, p=0.5, direction="horizontal", flip_idx=None) -> None:
         """
         Initializes the RandomFlip class with probability and direction.
 
+        This class applies a random horizontal or vertical flip to an image with a given probability.
+        It also updates any instances (bounding boxes, keypoints, etc.) accordingly.
+
         Args:
-            p (float, optional): The probability of applying the flip. Must be between 0 and 1. Default is 0.5.
-            direction (str, optional): The direction to apply the flip. Must be 'horizontal' or 'vertical'.
-                Default is 'horizontal'.
-            flip_idx (array-like, optional): Index mapping for flipping keypoints, if any.
+            p (float): The probability of applying the flip. Must be between 0 and 1.
+            direction (str): The direction to apply the flip. Must be 'horizontal' or 'vertical'.
+            flip_idx (List[int] | None): Index mapping for flipping keypoints, if any.
+
+        Raises:
+            AssertionError: If direction is not 'horizontal' or 'vertical', or if p is not between 0 and 1.
+
+        Examples:
+            >>> flip = RandomFlip(p=0.5, direction="horizontal")
+            >>> flip_with_idx = RandomFlip(p=0.7, direction="vertical", flip_idx=[1, 0, 3, 2, 5, 4])
         """
-        assert direction in ["horizontal", "vertical"], f"Support direction `horizontal` or `vertical`, got {direction}"
-        assert 0 <= p <= 1.0
+        assert direction in {"horizontal", "vertical"}, f"Support direction `horizontal` or `vertical`, got {direction}"
+        assert 0 <= p <= 1.0, f"The probability should be in range [0, 1], but got {p}."
 
         self.p = p
         self.direction = direction
@@ -652,12 +1432,25 @@ class RandomFlip:
         """
         Applies random flip to an image and updates any instances like bounding boxes or keypoints accordingly.
 
+        This method randomly flips the input image either horizontally or vertically based on the initialized
+        probability and direction. It also updates the corresponding instances (bounding boxes, keypoints) to
+        match the flipped image.
+
         Args:
-            labels (dict): A dictionary containing the keys 'img' and 'instances'. 'img' is the image to be flipped.
-                           'instances' is an object containing bounding boxes and optionally keypoints.
+            labels (Dict): A dictionary containing the following keys:
+                'img' (numpy.ndarray): The image to be flipped.
+                'instances' (ultralytics.utils.instance.Instances): An object containing bounding boxes and
+                    optionally keypoints.
 
         Returns:
-            (dict): The same dict with the flipped image and updated instances under the 'img' and 'instances' keys.
+            (Dict): The same dictionary with the flipped image and updated instances:
+                'img' (numpy.ndarray): The flipped image.
+                'instances' (ultralytics.utils.instance.Instances): Updated instances matching the flipped image.
+
+        Examples:
+            >>> labels = {"img": np.random.rand(640, 640, 3), "instances": Instances(...)}
+            >>> random_flip = RandomFlip(p=0.5, direction="horizontal")
+            >>> flipped_labels = random_flip(labels)
         """
         img = labels["img"]
         instances = labels.pop("instances")
@@ -682,10 +1475,56 @@ class RandomFlip:
 
 
 class LetterBox:
-    """Resize image and padding for detection, instance segmentation, pose."""
+    """
+    Resize image and padding for detection, instance segmentation, pose.
+
+    This class resizes and pads images to a specified shape while preserving aspect ratio. It also updates
+    corresponding labels and bounding boxes.
+
+    Attributes:
+        new_shape (tuple): Target shape (height, width) for resizing.
+        auto (bool): Whether to use minimum rectangle.
+        scaleFill (bool): Whether to stretch the image to new_shape.
+        scaleup (bool): Whether to allow scaling up. If False, only scale down.
+        stride (int): Stride for rounding padding.
+        center (bool): Whether to center the image or align to top-left.
+
+    Methods:
+        __call__: Resize and pad image, update labels and bounding boxes.
+
+    Examples:
+        >>> transform = LetterBox(new_shape=(640, 640))
+        >>> result = transform(labels)
+        >>> resized_img = result["img"]
+        >>> updated_instances = result["instances"]
+    """
 
     def __init__(self, new_shape=(640, 640), auto=False, scaleFill=False, scaleup=True, center=True, stride=32):
-        """Initialize LetterBox object with specific parameters."""
+        """
+        Initialize LetterBox object for resizing and padding images.
+
+        This class is designed to resize and pad images for object detection, instance segmentation, and pose estimation
+        tasks. It supports various resizing modes including auto-sizing, scale-fill, and letterboxing.
+
+        Args:
+            new_shape (Tuple[int, int]): Target size (height, width) for the resized image.
+            auto (bool): If True, use minimum rectangle to resize. If False, use new_shape directly.
+            scaleFill (bool): If True, stretch the image to new_shape without padding.
+            scaleup (bool): If True, allow scaling up. If False, only scale down.
+            center (bool): If True, center the placed image. If False, place image in top-left corner.
+            stride (int): Stride of the model (e.g., 32 for YOLOv5).
+
+        Attributes:
+            new_shape (Tuple[int, int]): Target size for the resized image.
+            auto (bool): Flag for using minimum rectangle resizing.
+            scaleFill (bool): Flag for stretching image without padding.
+            scaleup (bool): Flag for allowing upscaling.
+            stride (int): Stride value for ensuring image size is divisible by stride.
+
+        Examples:
+            >>> letterbox = LetterBox(new_shape=(640, 640), auto=False, scaleFill=False, scaleup=True, stride=32)
+            >>> resized_img = letterbox(original_img)
+        """
         self.new_shape = new_shape
         self.auto = auto
         self.scaleFill = scaleFill
@@ -694,7 +1533,27 @@ class LetterBox:
         self.center = center  # Put the image in the middle or top-left
 
     def __call__(self, labels=None, image=None):
-        """Return updated labels and image with added border."""
+        """
+        Resizes and pads an image for object detection, instance segmentation, or pose estimation tasks.
+
+        This method applies letterboxing to the input image, which involves resizing the image while maintaining its
+        aspect ratio and adding padding to fit the new shape. It also updates any associated labels accordingly.
+
+        Args:
+            labels (Dict | None): A dictionary containing image data and associated labels, or empty dict if None.
+            image (np.ndarray | None): The input image as a numpy array. If None, the image is taken from 'labels'.
+
+        Returns:
+            (Dict | Tuple): If 'labels' is provided, returns an updated dictionary with the resized and padded image,
+                updated labels, and additional metadata. If 'labels' is empty, returns a tuple containing the resized
+                and padded image, and a tuple of (ratio, (left_pad, top_pad)).
+
+        Examples:
+            >>> letterbox = LetterBox(new_shape=(640, 640))
+            >>> result = letterbox(labels={"img": np.zeros((480, 640, 3)), "instances": Instances(...)})
+            >>> resized_img = result["img"]
+            >>> updated_instances = result["instances"]
+        """
         if labels is None:
             labels = {}
         img = labels.get("img") if image is None else image
@@ -734,15 +1593,37 @@ class LetterBox:
             labels["ratio_pad"] = (labels["ratio_pad"], (left, top))  # for evaluation
 
         if len(labels):
-            labels = self._update_labels(labels, ratio, dw, dh)
+            labels = self._update_labels(labels, ratio, left, top)
             labels["img"] = img
             labels["resized_shape"] = new_shape
             return labels
         else:
             return img
 
-    def _update_labels(self, labels, ratio, padw, padh):
-        """Update labels."""
+    @staticmethod
+    def _update_labels(labels, ratio, padw, padh):
+        """
+        Updates labels after applying letterboxing to an image.
+
+        This method modifies the bounding box coordinates of instances in the labels
+        to account for resizing and padding applied during letterboxing.
+
+        Args:
+            labels (Dict): A dictionary containing image labels and instances.
+            ratio (Tuple[float, float]): Scaling ratios (width, height) applied to the image.
+            padw (float): Padding width added to the image.
+            padh (float): Padding height added to the image.
+
+        Returns:
+            (Dict): Updated labels dictionary with modified instance coordinates.
+
+        Examples:
+            >>> letterbox = LetterBox(new_shape=(640, 640))
+            >>> labels = {"instances": Instances(...)}
+            >>> ratio = (0.5, 0.5)
+            >>> padw, padh = 10, 20
+            >>> updated_labels = letterbox._update_labels(labels, ratio, padw, padh)
+        """
         labels["instances"].convert_bbox(format="xyxy")
         labels["instances"].denormalize(*labels["img"].shape[:2][::-1])
         labels["instances"].scale(*ratio)
@@ -750,91 +1631,217 @@ class LetterBox:
         return labels
 
 
-class CopyPaste:
+class CopyPaste(BaseMixTransform):
     """
-    Implements the Copy-Paste augmentation as described in the paper https://arxiv.org/abs/2012.07177. This class is
-    responsible for applying the Copy-Paste augmentation on images and their corresponding instances.
+    CopyPaste class for applying Copy-Paste augmentation to image datasets.
+
+    This class implements the Copy-Paste augmentation technique as described in the paper "Simple Copy-Paste is a Strong
+    Data Augmentation Method for Instance Segmentation" (https://arxiv.org/abs/2012.07177). It combines objects from
+    different images to create new training samples.
+
+    Attributes:
+        dataset (Any): The dataset to which Copy-Paste augmentation will be applied.
+        pre_transform (Callable | None): Optional transform to apply before Copy-Paste.
+        p (float): Probability of applying Copy-Paste augmentation.
+
+    Methods:
+        get_indexes: Returns a random index from the dataset.
+        _mix_transform: Applies Copy-Paste augmentation to the input labels.
+        __call__: Applies the Copy-Paste transformation to images and annotations.
+
+    Examples:
+        >>> from ultralytics.data.augment import CopyPaste
+        >>> dataset = YourDataset(...)  # Your image dataset
+        >>> copypaste = CopyPaste(dataset, p=0.5)
+        >>> augmented_labels = copypaste(original_labels)
     """
 
-    def __init__(self, p=0.5) -> None:
-        """
-        Initializes the CopyPaste class with a given probability.
+    def __init__(self, dataset=None, pre_transform=None, p=0.5, mode="flip") -> None:
+        """Initializes CopyPaste object with dataset, pre_transform, and probability of applying MixUp."""
+        super().__init__(dataset=dataset, pre_transform=pre_transform, p=p)
+        assert mode in {"flip", "mixup"}, f"Expected `mode` to be `flip` or `mixup`, but got {mode}."
+        self.mode = mode
 
-        Args:
-            p (float, optional): The probability of applying the Copy-Paste augmentation. Must be between 0 and 1.
-                                 Default is 0.5.
-        """
-        self.p = p
+    def get_indexes(self):
+        """Returns a list of random indexes from the dataset for CopyPaste augmentation."""
+        return random.randint(0, len(self.dataset) - 1)
+
+    def _mix_transform(self, labels):
+        """Applies Copy-Paste augmentation to combine objects from another image into the current image."""
+        labels2 = labels["mix_labels"][0]
+        return self._transform(labels, labels2)
 
     def __call__(self, labels):
-        """
-        Applies the Copy-Paste augmentation to the given image and instances.
+        """Applies Copy-Paste augmentation to an image and its labels."""
+        if len(labels["instances"].segments) == 0 or self.p == 0:
+            return labels
+        if self.mode == "flip":
+            return self._transform(labels)
 
-        Args:
-            labels (dict): A dictionary containing:
-                           - 'img': The image to augment.
-                           - 'cls': Class labels associated with the instances.
-                           - 'instances': Object containing bounding boxes, and optionally, keypoints and segments.
+        # Get index of one or three other images
+        indexes = self.get_indexes()
+        if isinstance(indexes, int):
+            indexes = [indexes]
 
-        Returns:
-            (dict): Dict with augmented image and updated instances under the 'img', 'cls', and 'instances' keys.
+        # Get images information will be used for Mosaic or MixUp
+        mix_labels = [self.dataset.get_image_and_label(i) for i in indexes]
 
-        Notes:
-            1. Instances are expected to have 'segments' as one of their attributes for this augmentation to work.
-            2. This method modifies the input dictionary 'labels' in place.
-        """
-        im = labels["img"]
-        cls = labels["cls"]
+        if self.pre_transform is not None:
+            for i, data in enumerate(mix_labels):
+                mix_labels[i] = self.pre_transform(data)
+        labels["mix_labels"] = mix_labels
+
+        # Update cls and texts
+        labels = self._update_label_text(labels)
+        # Mosaic or MixUp
+        labels = self._mix_transform(labels)
+        labels.pop("mix_labels", None)
+        return labels
+
+    def _transform(self, labels1, labels2={}):
+        """Applies Copy-Paste augmentation to combine objects from another image into the current image."""
+        im = labels1["img"]
+        cls = labels1["cls"]
         h, w = im.shape[:2]
-        instances = labels.pop("instances")
+        instances = labels1.pop("instances")
         instances.convert_bbox(format="xyxy")
         instances.denormalize(w, h)
-        if self.p and len(instances.segments):
-            n = len(instances)
-            _, w, _ = im.shape  # height, width, channels
-            im_new = np.zeros(im.shape, np.uint8)
-
-            # Calculate ioa first then select indexes randomly
-            ins_flip = deepcopy(instances)
-            ins_flip.fliplr(w)
-
-            ioa = bbox_ioa(ins_flip.bboxes, instances.bboxes)  # intersection over area, (N, M)
-            indexes = np.nonzero((ioa < 0.30).all(1))[0]  # (N, )
-            n = len(indexes)
-            for j in random.sample(list(indexes), k=round(self.p * n)):
-                cls = np.concatenate((cls, cls[[j]]), axis=0)
-                instances = Instances.concatenate((instances, ins_flip[[j]]), axis=0)
-                cv2.drawContours(im_new, instances.segments[[j]].astype(np.int32), -1, (1, 1, 1), cv2.FILLED)
-
-            result = cv2.flip(im, 1)  # augment segments (flip left-right)
-            i = cv2.flip(im_new, 1).astype(bool)
-            im[i] = result[i]
-
-        labels["img"] = im
-        labels["cls"] = cls
-        labels["instances"] = instances
-        return labels
+
+        im_new = np.zeros(im.shape, np.uint8)
+        instances2 = labels2.pop("instances", None)
+        if instances2 is None:
+            instances2 = deepcopy(instances)
+            instances2.fliplr(w)
+        ioa = bbox_ioa(instances2.bboxes, instances.bboxes)  # intersection over area, (N, M)
+        indexes = np.nonzero((ioa < 0.30).all(1))[0]  # (N, )
+        n = len(indexes)
+        sorted_idx = np.argsort(ioa.max(1)[indexes])
+        indexes = indexes[sorted_idx]
+        for j in indexes[: round(self.p * n)]:
+            cls = np.concatenate((cls, labels2.get("cls", cls)[[j]]), axis=0)
+            instances = Instances.concatenate((instances, instances2[[j]]), axis=0)
+            cv2.drawContours(im_new, instances2.segments[[j]].astype(np.int32), -1, (1, 1, 1), cv2.FILLED)
+
+        result = labels2.get("img", cv2.flip(im, 1))  # augment segments
+        i = im_new.astype(bool)
+        im[i] = result[i]
+
+        labels1["img"] = im
+        labels1["cls"] = cls
+        labels1["instances"] = instances
+        return labels1
 
 
 class Albumentations:
     """
-    Albumentations transformations.
+    Albumentations transformations for image augmentation.
+
+    This class applies various image transformations using the Albumentations library. It includes operations such as
+    Blur, Median Blur, conversion to grayscale, Contrast Limited Adaptive Histogram Equalization (CLAHE), random changes
+    in brightness and contrast, RandomGamma, and image quality reduction through compression.
 
-    Optional, uninstall package to disable. Applies Blur, Median Blur, convert to grayscale, Contrast Limited Adaptive
-    Histogram Equalization, random change of brightness and contrast, RandomGamma and lowering of image quality by
-    compression.
+    Attributes:
+        p (float): Probability of applying the transformations.
+        transform (albumentations.Compose): Composed Albumentations transforms.
+        contains_spatial (bool): Indicates if the transforms include spatial operations.
+
+    Methods:
+        __call__: Applies the Albumentations transformations to the input labels.
+
+    Examples:
+        >>> transform = Albumentations(p=0.5)
+        >>> augmented_labels = transform(labels)
+
+    Notes:
+        - The Albumentations package must be installed to use this class.
+        - If the package is not installed or an error occurs during initialization, the transform will be set to None.
+        - Spatial transforms are handled differently and require special processing for bounding boxes.
     """
 
     def __init__(self, p=1.0):
-        """Initialize the transform object for YOLO bbox formatted params."""
+        """
+        Initialize the Albumentations transform object for YOLO bbox formatted parameters.
+
+        This class applies various image augmentations using the Albumentations library, including Blur, Median Blur,
+        conversion to grayscale, Contrast Limited Adaptive Histogram Equalization, random changes of brightness and
+        contrast, RandomGamma, and image quality reduction through compression.
+
+        Args:
+            p (float): Probability of applying the augmentations. Must be between 0 and 1.
+
+        Attributes:
+            p (float): Probability of applying the augmentations.
+            transform (albumentations.Compose): Composed Albumentations transforms.
+            contains_spatial (bool): Indicates if the transforms include spatial transformations.
+
+        Raises:
+            ImportError: If the Albumentations package is not installed.
+            Exception: For any other errors during initialization.
+
+        Examples:
+            >>> transform = Albumentations(p=0.5)
+            >>> augmented = transform(image=image, bboxes=bboxes, class_labels=classes)
+            >>> augmented_image = augmented["image"]
+            >>> augmented_bboxes = augmented["bboxes"]
+
+        Notes:
+            - Requires Albumentations version 1.0.3 or higher.
+            - Spatial transforms are handled differently to ensure bbox compatibility.
+            - Some transforms are applied with very low probability (0.01) by default.
+        """
         self.p = p
         self.transform = None
         prefix = colorstr("albumentations: ")
+
         try:
             import albumentations as A
 
             check_version(A.__version__, "1.0.3", hard=True)  # version requirement
 
+            # List of possible spatial transforms
+            spatial_transforms = {
+                "Affine",
+                "BBoxSafeRandomCrop",
+                "CenterCrop",
+                "CoarseDropout",
+                "Crop",
+                "CropAndPad",
+                "CropNonEmptyMaskIfExists",
+                "D4",
+                "ElasticTransform",
+                "Flip",
+                "GridDistortion",
+                "GridDropout",
+                "HorizontalFlip",
+                "Lambda",
+                "LongestMaxSize",
+                "MaskDropout",
+                "MixUp",
+                "Morphological",
+                "NoOp",
+                "OpticalDistortion",
+                "PadIfNeeded",
+                "Perspective",
+                "PiecewiseAffine",
+                "PixelDropout",
+                "RandomCrop",
+                "RandomCropFromBorders",
+                "RandomGridShuffle",
+                "RandomResizedCrop",
+                "RandomRotate90",
+                "RandomScale",
+                "RandomSizedBBoxSafeCrop",
+                "RandomSizedCrop",
+                "Resize",
+                "Rotate",
+                "SafeRotate",
+                "ShiftScaleRotate",
+                "SmallestMaxSize",
+                "Transpose",
+                "VerticalFlip",
+                "XYMasking",
+            }  # from https://albumentations.ai/docs/getting_started/transforms_and_targets/#spatial-level-transforms
+
             # Transforms
             T = [
                 A.Blur(p=0.01),
@@ -845,8 +1852,17 @@ class Albumentations:
                 A.RandomGamma(p=0.0),
                 A.ImageCompression(quality_lower=75, p=0.0),
             ]
-            self.transform = A.Compose(T, bbox_params=A.BboxParams(format="yolo", label_fields=["class_labels"]))
 
+            # Compose transforms
+            self.contains_spatial = any(transform.__class__.__name__ in spatial_transforms for transform in T)
+            self.transform = (
+                A.Compose(T, bbox_params=A.BboxParams(format="yolo", label_fields=["class_labels"]))
+                if self.contains_spatial
+                else A.Compose(T)
+            )
+            if hasattr(self.transform, "set_random_seed"):
+                # Required for deterministic transforms in albumentations>=1.4.21
+                self.transform.set_random_seed(torch.initial_seed())
             LOGGER.info(prefix + ", ".join(f"{x}".replace("always_apply=False, ", "") for x in T if x.p))
         except ImportError:  # package not installed, skip
             pass
@@ -854,38 +1870,87 @@ class Albumentations:
             LOGGER.info(f"{prefix}{e}")
 
     def __call__(self, labels):
-        """Generates object detections and returns a dictionary with detection results."""
-        im = labels["img"]
-        cls = labels["cls"]
-        if len(cls):
-            labels["instances"].convert_bbox("xywh")
-            labels["instances"].normalize(*im.shape[:2][::-1])
-            bboxes = labels["instances"].bboxes
-            # TODO: add supports of segments and keypoints
-            if self.transform and random.random() < self.p:
+        """
+        Applies Albumentations transformations to input labels.
+
+        This method applies a series of image augmentations using the Albumentations library. It can perform both
+        spatial and non-spatial transformations on the input image and its corresponding labels.
+
+        Args:
+            labels (Dict): A dictionary containing image data and annotations. Expected keys are:
+                - 'img': numpy.ndarray representing the image
+                - 'cls': numpy.ndarray of class labels
+                - 'instances': object containing bounding boxes and other instance information
+
+        Returns:
+            (Dict): The input dictionary with augmented image and updated annotations.
+
+        Examples:
+            >>> transform = Albumentations(p=0.5)
+            >>> labels = {
+            ...     "img": np.random.rand(640, 640, 3),
+            ...     "cls": np.array([0, 1]),
+            ...     "instances": Instances(bboxes=np.array([[0, 0, 1, 1], [0.5, 0.5, 0.8, 0.8]])),
+            ... }
+            >>> augmented = transform(labels)
+            >>> assert augmented["img"].shape == (640, 640, 3)
+
+        Notes:
+            - The method applies transformations with probability self.p.
+            - Spatial transforms update bounding boxes, while non-spatial transforms only modify the image.
+            - Requires the Albumentations library to be installed.
+        """
+        if self.transform is None or random.random() > self.p:
+            return labels
+
+        if self.contains_spatial:
+            cls = labels["cls"]
+            if len(cls):
+                im = labels["img"]
+                labels["instances"].convert_bbox("xywh")
+                labels["instances"].normalize(*im.shape[:2][::-1])
+                bboxes = labels["instances"].bboxes
+                # TODO: add supports of segments and keypoints
                 new = self.transform(image=im, bboxes=bboxes, class_labels=cls)  # transformed
                 if len(new["class_labels"]) > 0:  # skip update if no bbox in new im
                     labels["img"] = new["image"]
                     labels["cls"] = np.array(new["class_labels"])
                     bboxes = np.array(new["bboxes"], dtype=np.float32)
-            labels["instances"].update(bboxes=bboxes)
+                labels["instances"].update(bboxes=bboxes)
+        else:
+            labels["img"] = self.transform(image=labels["img"])["image"]  # transformed
+
         return labels
 
 
-# TODO: technically this is not an augmentation, maybe we should put this to another files
 class Format:
     """
-    Formats image annotations for object detection, instance segmentation, and pose estimation tasks. The class
-    standardizes the image and instance annotations to be used by the `collate_fn` in PyTorch DataLoader.
+    A class for formatting image annotations for object detection, instance segmentation, and pose estimation tasks.
+
+    This class standardizes image and instance annotations to be used by the `collate_fn` in PyTorch DataLoader.
 
     Attributes:
-        bbox_format (str): Format for bounding boxes. Default is 'xywh'.
-        normalize (bool): Whether to normalize bounding boxes. Default is True.
-        return_mask (bool): Return instance masks for segmentation. Default is False.
-        return_keypoint (bool): Return keypoints for pose estimation. Default is False.
-        mask_ratio (int): Downsample ratio for masks. Default is 4.
-        mask_overlap (bool): Whether to overlap masks. Default is True.
-        batch_idx (bool): Keep batch indexes. Default is True.
+        bbox_format (str): Format for bounding boxes. Options are 'xywh' or 'xyxy'.
+        normalize (bool): Whether to normalize bounding boxes.
+        return_mask (bool): Whether to return instance masks for segmentation.
+        return_keypoint (bool): Whether to return keypoints for pose estimation.
+        return_obb (bool): Whether to return oriented bounding boxes.
+        mask_ratio (int): Downsample ratio for masks.
+        mask_overlap (bool): Whether to overlap masks.
+        batch_idx (bool): Whether to keep batch indexes.
+        bgr (float): The probability to return BGR images.
+
+    Methods:
+        __call__: Formats labels dictionary with image, classes, bounding boxes, and optionally masks and keypoints.
+        _format_img: Converts image from Numpy array to PyTorch tensor.
+        _format_segments: Converts polygon points to bitmap masks.
+
+    Examples:
+        >>> formatter = Format(bbox_format="xywh", normalize=True, return_mask=True)
+        >>> formatted_labels = formatter(labels)
+        >>> img = formatted_labels["img"]
+        >>> bboxes = formatted_labels["bboxes"]
+        >>> masks = formatted_labels["masks"]
     """
 
     def __init__(
@@ -898,8 +1963,41 @@ class Format:
         mask_ratio=4,
         mask_overlap=True,
         batch_idx=True,
+        bgr=0.0,
     ):
-        """Initializes the Format class with given parameters."""
+        """
+        Initializes the Format class with given parameters for image and instance annotation formatting.
+
+        This class standardizes image and instance annotations for object detection, instance segmentation, and pose
+        estimation tasks, preparing them for use in PyTorch DataLoader's `collate_fn`.
+
+        Args:
+            bbox_format (str): Format for bounding boxes. Options are 'xywh', 'xyxy', etc.
+            normalize (bool): Whether to normalize bounding boxes to [0,1].
+            return_mask (bool): If True, returns instance masks for segmentation tasks.
+            return_keypoint (bool): If True, returns keypoints for pose estimation tasks.
+            return_obb (bool): If True, returns oriented bounding boxes.
+            mask_ratio (int): Downsample ratio for masks.
+            mask_overlap (bool): If True, allows mask overlap.
+            batch_idx (bool): If True, keeps batch indexes.
+            bgr (float): Probability of returning BGR images instead of RGB.
+
+        Attributes:
+            bbox_format (str): Format for bounding boxes.
+            normalize (bool): Whether bounding boxes are normalized.
+            return_mask (bool): Whether to return instance masks.
+            return_keypoint (bool): Whether to return keypoints.
+            return_obb (bool): Whether to return oriented bounding boxes.
+            mask_ratio (int): Downsample ratio for masks.
+            mask_overlap (bool): Whether masks can overlap.
+            batch_idx (bool): Whether to keep batch indexes.
+            bgr (float): The probability to return BGR images.
+
+        Examples:
+            >>> format = Format(bbox_format="xyxy", return_mask=True, return_keypoint=False)
+            >>> print(format.bbox_format)
+            xyxy
+        """
         self.bbox_format = bbox_format
         self.normalize = normalize
         self.return_mask = return_mask  # set False when training detection only
@@ -908,9 +2006,37 @@ class Format:
         self.mask_ratio = mask_ratio
         self.mask_overlap = mask_overlap
         self.batch_idx = batch_idx  # keep the batch indexes
+        self.bgr = bgr
 
     def __call__(self, labels):
-        """Return formatted image, classes, bounding boxes & keypoints to be used by 'collate_fn'."""
+        """
+        Formats image annotations for object detection, instance segmentation, and pose estimation tasks.
+
+        This method standardizes the image and instance annotations to be used by the `collate_fn` in PyTorch
+        DataLoader. It processes the input labels dictionary, converting annotations to the specified format and
+        applying normalization if required.
+
+        Args:
+            labels (Dict): A dictionary containing image and annotation data with the following keys:
+                - 'img': The input image as a numpy array.
+                - 'cls': Class labels for instances.
+                - 'instances': An Instances object containing bounding boxes, segments, and keypoints.
+
+        Returns:
+            (Dict): A dictionary with formatted data, including:
+                - 'img': Formatted image tensor.
+                - 'cls': Class label's tensor.
+                - 'bboxes': Bounding boxes tensor in the specified format.
+                - 'masks': Instance masks tensor (if return_mask is True).
+                - 'keypoints': Keypoints tensor (if return_keypoint is True).
+                - 'batch_idx': Batch index tensor (if batch_idx is True).
+
+        Examples:
+            >>> formatter = Format(bbox_format="xywh", normalize=True, return_mask=True)
+            >>> labels = {"img": np.random.rand(640, 640, 3), "cls": np.array([0, 1]), "instances": Instances(...)}
+            >>> formatted_labels = formatter(labels)
+            >>> print(formatted_labels.keys())
+        """
         img = labels.pop("img")
         h, w = img.shape[:2]
         cls = labels.pop("cls")
@@ -928,32 +2054,78 @@ class Format:
                     1 if self.mask_overlap else nl, img.shape[0] // self.mask_ratio, img.shape[1] // self.mask_ratio
                 )
             labels["masks"] = masks
-        if self.normalize:
-            instances.normalize(w, h)
         labels["img"] = self._format_img(img)
         labels["cls"] = torch.from_numpy(cls) if nl else torch.zeros(nl)
         labels["bboxes"] = torch.from_numpy(instances.bboxes) if nl else torch.zeros((nl, 4))
         if self.return_keypoint:
             labels["keypoints"] = torch.from_numpy(instances.keypoints)
+            if self.normalize:
+                labels["keypoints"][..., 0] /= w
+                labels["keypoints"][..., 1] /= h
         if self.return_obb:
             labels["bboxes"] = (
                 xyxyxyxy2xywhr(torch.from_numpy(instances.segments)) if len(instances.segments) else torch.zeros((0, 5))
             )
+        # NOTE: need to normalize obb in xywhr format for width-height consistency
+        if self.normalize:
+            labels["bboxes"][:, [0, 2]] /= w
+            labels["bboxes"][:, [1, 3]] /= h
         # Then we can use collate_fn
         if self.batch_idx:
             labels["batch_idx"] = torch.zeros(nl)
         return labels
 
     def _format_img(self, img):
-        """Format the image for YOLO from Numpy array to PyTorch tensor."""
+        """
+        Formats an image for YOLO from a Numpy array to a PyTorch tensor.
+
+        This function performs the following operations:
+        1. Ensures the image has 3 dimensions (adds a channel dimension if needed).
+        2. Transposes the image from HWC to CHW format.
+        3. Optionally flips the color channels from RGB to BGR.
+        4. Converts the image to a contiguous array.
+        5. Converts the Numpy array to a PyTorch tensor.
+
+        Args:
+            img (np.ndarray): Input image as a Numpy array with shape (H, W, C) or (H, W).
+
+        Returns:
+            (torch.Tensor): Formatted image as a PyTorch tensor with shape (C, H, W).
+
+        Examples:
+            >>> import numpy as np
+            >>> img = np.random.rand(100, 100, 3)
+            >>> formatted_img = self._format_img(img)
+            >>> print(formatted_img.shape)
+            torch.Size([3, 100, 100])
+        """
         if len(img.shape) < 3:
             img = np.expand_dims(img, -1)
-        img = np.ascontiguousarray(img.transpose(2, 0, 1)[::-1])
+        img = img.transpose(2, 0, 1)
+        img = np.ascontiguousarray(img[::-1] if random.uniform(0, 1) > self.bgr else img)
         img = torch.from_numpy(img)
         return img
 
     def _format_segments(self, instances, cls, w, h):
-        """Convert polygon points to bitmap."""
+        """
+        Converts polygon segments to bitmap masks.
+
+        Args:
+            instances (Instances): Object containing segment information.
+            cls (numpy.ndarray): Class labels for each instance.
+            w (int): Width of the image.
+            h (int): Height of the image.
+
+        Returns:
+            masks (numpy.ndarray): Bitmap masks with shape (N, H, W) or (1, H, W) if mask_overlap is True.
+            instances (Instances): Updated instances object with sorted segments if mask_overlap is True.
+            cls (numpy.ndarray): Updated class labels, sorted if mask_overlap is True.
+
+        Notes:
+            - If self.mask_overlap is True, masks are overlapped and sorted by area.
+            - If self.mask_overlap is False, each mask is represented separately.
+            - Masks are downsampled according to self.mask_ratio.
+        """
         segments = instances.segments
         if self.mask_overlap:
             masks, sorted_idx = polygons2masks_overlap((h, w), segments, downsample_ratio=self.mask_ratio)
@@ -966,22 +2138,189 @@ class Format:
         return masks, instances, cls
 
 
+class RandomLoadText:
+    """
+    Randomly samples positive and negative texts and updates class indices accordingly.
+
+    This class is responsible for sampling texts from a given set of class texts, including both positive
+    (present in the image) and negative (not present in the image) samples. It updates the class indices
+    to reflect the sampled texts and can optionally pad the text list to a fixed length.
+
+    Attributes:
+        prompt_format (str): Format string for text prompts.
+        neg_samples (Tuple[int, int]): Range for randomly sampling negative texts.
+        max_samples (int): Maximum number of different text samples in one image.
+        padding (bool): Whether to pad texts to max_samples.
+        padding_value (str): The text used for padding when padding is True.
+
+    Methods:
+        __call__: Processes the input labels and returns updated classes and texts.
+
+    Examples:
+        >>> loader = RandomLoadText(prompt_format="Object: {}", neg_samples=(5, 10), max_samples=20)
+        >>> labels = {"cls": [0, 1, 2], "texts": [["cat"], ["dog"], ["bird"]], "instances": [...]}
+        >>> updated_labels = loader(labels)
+        >>> print(updated_labels["texts"])
+        ['Object: cat', 'Object: dog', 'Object: bird', 'Object: elephant', 'Object: car']
+    """
+
+    def __init__(
+        self,
+        prompt_format: str = "{}",
+        neg_samples: Tuple[int, int] = (80, 80),
+        max_samples: int = 80,
+        padding: bool = False,
+        padding_value: str = "",
+    ) -> None:
+        """
+        Initializes the RandomLoadText class for randomly sampling positive and negative texts.
+
+        This class is designed to randomly sample positive texts and negative texts, and update the class
+        indices accordingly to the number of samples. It can be used for text-based object detection tasks.
+
+        Args:
+            prompt_format (str): Format string for the prompt. Default is '{}'. The format string should
+                contain a single pair of curly braces {} where the text will be inserted.
+            neg_samples (Tuple[int, int]): A range to randomly sample negative texts. The first integer
+                specifies the minimum number of negative samples, and the second integer specifies the
+                maximum. Default is (80, 80).
+            max_samples (int): The maximum number of different text samples in one image. Default is 80.
+            padding (bool): Whether to pad texts to max_samples. If True, the number of texts will always
+                be equal to max_samples. Default is False.
+            padding_value (str): The padding text to use when padding is True. Default is an empty string.
+
+        Attributes:
+            prompt_format (str): The format string for the prompt.
+            neg_samples (Tuple[int, int]): The range for sampling negative texts.
+            max_samples (int): The maximum number of text samples.
+            padding (bool): Whether padding is enabled.
+            padding_value (str): The value used for padding.
+
+        Examples:
+            >>> random_load_text = RandomLoadText(prompt_format="Object: {}", neg_samples=(50, 100), max_samples=120)
+            >>> random_load_text.prompt_format
+            'Object: {}'
+            >>> random_load_text.neg_samples
+            (50, 100)
+            >>> random_load_text.max_samples
+            120
+        """
+        self.prompt_format = prompt_format
+        self.neg_samples = neg_samples
+        self.max_samples = max_samples
+        self.padding = padding
+        self.padding_value = padding_value
+
+    def __call__(self, labels: dict) -> dict:
+        """
+        Randomly samples positive and negative texts and updates class indices accordingly.
+
+        This method samples positive texts based on the existing class labels in the image, and randomly
+        selects negative texts from the remaining classes. It then updates the class indices to match the
+        new sampled text order.
+
+        Args:
+            labels (Dict): A dictionary containing image labels and metadata. Must include 'texts' and 'cls' keys.
+
+        Returns:
+            (Dict): Updated labels dictionary with new 'cls' and 'texts' entries.
+
+        Examples:
+            >>> loader = RandomLoadText(prompt_format="A photo of {}", neg_samples=(5, 10), max_samples=20)
+            >>> labels = {"cls": np.array([[0], [1], [2]]), "texts": [["dog"], ["cat"], ["bird"]]}
+            >>> updated_labels = loader(labels)
+        """
+        assert "texts" in labels, "No texts found in labels."
+        class_texts = labels["texts"]
+        num_classes = len(class_texts)
+        cls = np.asarray(labels.pop("cls"), dtype=int)
+        pos_labels = np.unique(cls).tolist()
+
+        if len(pos_labels) > self.max_samples:
+            pos_labels = random.sample(pos_labels, k=self.max_samples)
+
+        neg_samples = min(min(num_classes, self.max_samples) - len(pos_labels), random.randint(*self.neg_samples))
+        neg_labels = [i for i in range(num_classes) if i not in pos_labels]
+        neg_labels = random.sample(neg_labels, k=neg_samples)
+
+        sampled_labels = pos_labels + neg_labels
+        random.shuffle(sampled_labels)
+
+        label2ids = {label: i for i, label in enumerate(sampled_labels)}
+        valid_idx = np.zeros(len(labels["instances"]), dtype=bool)
+        new_cls = []
+        for i, label in enumerate(cls.squeeze(-1).tolist()):
+            if label not in label2ids:
+                continue
+            valid_idx[i] = True
+            new_cls.append([label2ids[label]])
+        labels["instances"] = labels["instances"][valid_idx]
+        labels["cls"] = np.array(new_cls)
+
+        # Randomly select one prompt when there's more than one prompts
+        texts = []
+        for label in sampled_labels:
+            prompts = class_texts[label]
+            assert len(prompts) > 0
+            prompt = self.prompt_format.format(prompts[random.randrange(len(prompts))])
+            texts.append(prompt)
+
+        if self.padding:
+            valid_labels = len(pos_labels) + len(neg_labels)
+            num_padding = self.max_samples - valid_labels
+            if num_padding > 0:
+                texts += [self.padding_value] * num_padding
+
+        labels["texts"] = texts
+        return labels
+
+
 def v8_transforms(dataset, imgsz, hyp, stretch=False):
-    """Convert images to a size suitable for YOLOv8 training."""
-    pre_transform = Compose(
-        [
-            Mosaic(dataset, imgsz=imgsz, p=hyp.mosaic),
-            CopyPaste(p=hyp.copy_paste),
-            RandomPerspective(
-                degrees=hyp.degrees,
-                translate=hyp.translate,
-                scale=hyp.scale,
-                shear=hyp.shear,
-                perspective=hyp.perspective,
-                pre_transform=None if stretch else LetterBox(new_shape=(imgsz, imgsz)),
-            ),
-        ]
+    """
+    Applies a series of image transformations for training.
+
+    This function creates a composition of image augmentation techniques to prepare images for YOLO training.
+    It includes operations such as mosaic, copy-paste, random perspective, mixup, and various color adjustments.
+
+    Args:
+        dataset (Dataset): The dataset object containing image data and annotations.
+        imgsz (int): The target image size for resizing.
+        hyp (Namespace): A dictionary of hyperparameters controlling various aspects of the transformations.
+        stretch (bool): If True, applies stretching to the image. If False, uses LetterBox resizing.
+
+    Returns:
+        (Compose): A composition of image transformations to be applied to the dataset.
+
+    Examples:
+        >>> from ultralytics.data.dataset import YOLODataset
+        >>> from ultralytics.utils import IterableSimpleNamespace
+        >>> dataset = YOLODataset(img_path="path/to/images", imgsz=640)
+        >>> hyp = IterableSimpleNamespace(mosaic=1.0, copy_paste=0.5, degrees=10.0, translate=0.2, scale=0.9)
+        >>> transforms = v8_transforms(dataset, imgsz=640, hyp=hyp)
+        >>> augmented_data = transforms(dataset[0])
+    """
+    mosaic = Mosaic(dataset, imgsz=imgsz, p=hyp.mosaic)
+    affine = RandomPerspective(
+        degrees=hyp.degrees,
+        translate=hyp.translate,
+        scale=hyp.scale,
+        shear=hyp.shear,
+        perspective=hyp.perspective,
+        pre_transform=None if stretch else LetterBox(new_shape=(imgsz, imgsz)),
     )
+
+    pre_transform = Compose([mosaic, affine])
+    if hyp.copy_paste_mode == "flip":
+        pre_transform.insert(1, CopyPaste(p=hyp.copy_paste, mode=hyp.copy_paste_mode))
+    else:
+        pre_transform.append(
+            CopyPaste(
+                dataset,
+                pre_transform=Compose([Mosaic(dataset, imgsz=imgsz, p=hyp.mosaic), affine]),
+                p=hyp.copy_paste,
+                mode=hyp.copy_paste_mode,
+            )
+        )
     flip_idx = dataset.data.get("flip_idx", [])  # for keypoints augmentation
     if dataset.use_keypoints:
         kpt_shape = dataset.data.get("kpt_shape", None)
@@ -1008,51 +2347,59 @@ def classify_transforms(
     size=224,
     mean=DEFAULT_MEAN,
     std=DEFAULT_STD,
-    interpolation: T.InterpolationMode = T.InterpolationMode.BILINEAR,
-    crop_fraction: float = DEFAULT_CROP_FTACTION,
+    interpolation="BILINEAR",
+    crop_fraction: float = DEFAULT_CROP_FRACTION,
 ):
     """
-    Classification transforms for evaluation/inference. Inspired by timm/data/transforms_factory.py.
+    Creates a composition of image transforms for classification tasks.
+
+    This function generates a sequence of torchvision transforms suitable for preprocessing images
+    for classification models during evaluation or inference. The transforms include resizing,
+    center cropping, conversion to tensor, and normalization.
 
     Args:
-        size (int): image size
-        mean (tuple): mean values of RGB channels
-        std (tuple): std values of RGB channels
-        interpolation (T.InterpolationMode): interpolation mode. default is T.InterpolationMode.BILINEAR.
-        crop_fraction (float): fraction of image to crop. default is 1.0.
+        size (int | tuple): The target size for the transformed image. If an int, it defines the shortest edge. If a
+            tuple, it defines (height, width).
+        mean (tuple): Mean values for each RGB channel used in normalization.
+        std (tuple): Standard deviation values for each RGB channel used in normalization.
+        interpolation (str): Interpolation method of either 'NEAREST', 'BILINEAR' or 'BICUBIC'.
+        crop_fraction (float): Fraction of the image to be cropped.
 
     Returns:
-        (T.Compose): torchvision transforms
+        (torchvision.transforms.Compose): A composition of torchvision transforms.
+
+    Examples:
+        >>> transforms = classify_transforms(size=224)
+        >>> img = Image.open("path/to/image.jpg")
+        >>> transformed_img = transforms(img)
     """
+    import torchvision.transforms as T  # scope for faster 'import ultralytics'
 
     if isinstance(size, (tuple, list)):
-        assert len(size) == 2
+        assert len(size) == 2, f"'size' tuples must be length 2, not length {len(size)}"
         scale_size = tuple(math.floor(x / crop_fraction) for x in size)
     else:
         scale_size = math.floor(size / crop_fraction)
         scale_size = (scale_size, scale_size)
 
-    # aspect ratio is preserved, crops center within image, no borders are added, image is lost
+    # Aspect ratio is preserved, crops center within image, no borders are added, image is lost
     if scale_size[0] == scale_size[1]:
-        # simple case, use torchvision built-in Resize w/ shortest edge mode (scalar size arg)
-        tfl = [T.Resize(scale_size[0], interpolation=interpolation)]
+        # Simple case, use torchvision built-in Resize with the shortest edge mode (scalar size arg)
+        tfl = [T.Resize(scale_size[0], interpolation=getattr(T.InterpolationMode, interpolation))]
     else:
-        # resize shortest edge to matching target dim for non-square target
+        # Resize the shortest edge to matching target dim for non-square target
         tfl = [T.Resize(scale_size)]
-    tfl += [T.CenterCrop(size)]
-
-    tfl += [
-        T.ToTensor(),
-        T.Normalize(
-            mean=torch.tensor(mean),
-            std=torch.tensor(std),
-        ),
-    ]
-
+    tfl.extend(
+        [
+            T.CenterCrop(size),
+            T.ToTensor(),
+            T.Normalize(mean=torch.tensor(mean), std=torch.tensor(std)),
+        ]
+    )
     return T.Compose(tfl)
 
 
-# Classification augmentations train ---------------------------------------------------------------------------------------
+# Classification training augmentations --------------------------------------------------------------------------------
 def classify_augmentations(
     size=224,
     mean=DEFAULT_MEAN,
@@ -1067,64 +2414,74 @@ def classify_augmentations(
     hsv_v=0.4,  # image HSV-Value augmentation (fraction)
     force_color_jitter=False,
     erasing=0.0,
-    interpolation: T.InterpolationMode = T.InterpolationMode.BILINEAR,
+    interpolation="BILINEAR",
 ):
     """
-    Classification transforms with augmentation for training. Inspired by timm/data/transforms_factory.py.
+    Creates a composition of image augmentation transforms for classification tasks.
+
+    This function generates a set of image transformations suitable for training classification models. It includes
+    options for resizing, flipping, color jittering, auto augmentation, and random erasing.
 
     Args:
-        size (int): image size
-        scale (tuple): scale range of the image. default is (0.08, 1.0)
-        ratio (tuple): aspect ratio range of the image. default is (3./4., 4./3.)
-        mean (tuple): mean values of RGB channels
-        std (tuple): std values of RGB channels
-        hflip (float): probability of horizontal flip
-        vflip (float): probability of vertical flip
-        auto_augment (str): auto augmentation policy. can be 'randaugment', 'augmix', 'autoaugment' or None.
-        hsv_h (float): image HSV-Hue augmentation (fraction)
-        hsv_s (float): image HSV-Saturation augmentation (fraction)
-        hsv_v (float): image HSV-Value augmentation (fraction)
-        force_color_jitter (bool): force to apply color jitter even if auto augment is enabled
-        erasing (float): probability of random erasing
-        interpolation (T.InterpolationMode): interpolation mode. default is T.InterpolationMode.BILINEAR.
+        size (int): Target size for the image after transformations.
+        mean (tuple): Mean values for normalization, one per channel.
+        std (tuple): Standard deviation values for normalization, one per channel.
+        scale (tuple | None): Range of size of the origin size cropped.
+        ratio (tuple | None): Range of aspect ratio of the origin aspect ratio cropped.
+        hflip (float): Probability of horizontal flip.
+        vflip (float): Probability of vertical flip.
+        auto_augment (str | None): Auto augmentation policy. Can be 'randaugment', 'augmix', 'autoaugment' or None.
+        hsv_h (float): Image HSV-Hue augmentation factor.
+        hsv_s (float): Image HSV-Saturation augmentation factor.
+        hsv_v (float): Image HSV-Value augmentation factor.
+        force_color_jitter (bool): Whether to apply color jitter even if auto augment is enabled.
+        erasing (float): Probability of random erasing.
+        interpolation (str): Interpolation method of either 'NEAREST', 'BILINEAR' or 'BICUBIC'.
 
     Returns:
-        (T.Compose): torchvision transforms
+        (torchvision.transforms.Compose): A composition of image augmentation transforms.
+
+    Examples:
+        >>> transforms = classify_augmentations(size=224, auto_augment="randaugment")
+        >>> augmented_image = transforms(original_image)
     """
-    # Transforms to apply if albumentations not installed
+    # Transforms to apply if Albumentations not installed
+    import torchvision.transforms as T  # scope for faster 'import ultralytics'
+
     if not isinstance(size, int):
         raise TypeError(f"classify_transforms() size {size} must be integer, not (list, tuple)")
     scale = tuple(scale or (0.08, 1.0))  # default imagenet scale range
     ratio = tuple(ratio or (3.0 / 4.0, 4.0 / 3.0))  # default imagenet ratio range
+    interpolation = getattr(T.InterpolationMode, interpolation)
     primary_tfl = [T.RandomResizedCrop(size, scale=scale, ratio=ratio, interpolation=interpolation)]
     if hflip > 0.0:
-        primary_tfl += [T.RandomHorizontalFlip(p=hflip)]
+        primary_tfl.append(T.RandomHorizontalFlip(p=hflip))
     if vflip > 0.0:
-        primary_tfl += [T.RandomVerticalFlip(p=vflip)]
+        primary_tfl.append(T.RandomVerticalFlip(p=vflip))
 
     secondary_tfl = []
     disable_color_jitter = False
     if auto_augment:
-        assert isinstance(auto_augment, str)
+        assert isinstance(auto_augment, str), f"Provided argument should be string, but got type {type(auto_augment)}"
         # color jitter is typically disabled if AA/RA on,
         # this allows override without breaking old hparm cfgs
         disable_color_jitter = not force_color_jitter
 
         if auto_augment == "randaugment":
             if TORCHVISION_0_11:
-                secondary_tfl += [T.RandAugment(interpolation=interpolation)]
+                secondary_tfl.append(T.RandAugment(interpolation=interpolation))
             else:
                 LOGGER.warning('"auto_augment=randaugment" requires torchvision >= 0.11.0. Disabling it.')
 
         elif auto_augment == "augmix":
             if TORCHVISION_0_13:
-                secondary_tfl += [T.AugMix(interpolation=interpolation)]
+                secondary_tfl.append(T.AugMix(interpolation=interpolation))
             else:
                 LOGGER.warning('"auto_augment=augmix" requires torchvision >= 0.13.0. Disabling it.')
 
         elif auto_augment == "autoaugment":
             if TORCHVISION_0_10:
-                secondary_tfl += [T.AutoAugment(interpolation=interpolation)]
+                secondary_tfl.append(T.AutoAugment(interpolation=interpolation))
             else:
                 LOGGER.warning('"auto_augment=autoaugment" requires torchvision >= 0.10.0. Disabling it.')
 
@@ -1135,7 +2492,7 @@ def classify_augmentations(
             )
 
     if not disable_color_jitter:
-        secondary_tfl += [T.ColorJitter(brightness=hsv_v, contrast=hsv_v, saturation=hsv_s, hue=hsv_h)]
+        secondary_tfl.append(T.ColorJitter(brightness=hsv_v, contrast=hsv_v, saturation=hsv_s, hue=hsv_h))
 
     final_tfl = [
         T.ToTensor(),
@@ -1149,24 +2506,53 @@ def classify_augmentations(
 # NOTE: keep this class for backward compatibility
 class ClassifyLetterBox:
     """
-    YOLOv8 LetterBox class for image preprocessing, designed to be part of a transformation pipeline, e.g.,
-    T.Compose([LetterBox(size), ToTensor()]).
+    A class for resizing and padding images for classification tasks.
+
+    This class is designed to be part of a transformation pipeline, e.g., T.Compose([LetterBox(size), ToTensor()]).
+    It resizes and pads images to a specified size while maintaining the original aspect ratio.
 
     Attributes:
         h (int): Target height of the image.
         w (int): Target width of the image.
-        auto (bool): If True, automatically solves for short side using stride.
+        auto (bool): If True, automatically calculates the short side using stride.
         stride (int): The stride value, used when 'auto' is True.
+
+    Methods:
+        __call__: Applies the letterbox transformation to an input image.
+
+    Examples:
+        >>> transform = ClassifyLetterBox(size=(640, 640), auto=False, stride=32)
+        >>> img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
+        >>> result = transform(img)
+        >>> print(result.shape)
+        (640, 640, 3)
     """
 
     def __init__(self, size=(640, 640), auto=False, stride=32):
         """
-        Initializes the ClassifyLetterBox class with a target size, auto-flag, and stride.
+        Initializes the ClassifyLetterBox object for image preprocessing.
+
+        This class is designed to be part of a transformation pipeline for image classification tasks. It resizes and
+        pads images to a specified size while maintaining the original aspect ratio.
 
         Args:
-            size (Union[int, Tuple[int, int]]): The target dimensions (height, width) for the letterbox.
-            auto (bool): If True, automatically calculates the short side based on stride.
-            stride (int): The stride value, used when 'auto' is True.
+            size (int | Tuple[int, int]): Target size for the letterboxed image. If an int, a square image of
+                (size, size) is created. If a tuple, it should be (height, width).
+            auto (bool): If True, automatically calculates the short side based on stride. Default is False.
+            stride (int): The stride value, used when 'auto' is True. Default is 32.
+
+        Attributes:
+            h (int): Target height of the letterboxed image.
+            w (int): Target width of the letterboxed image.
+            auto (bool): Flag indicating whether to automatically calculate short side.
+            stride (int): Stride value for automatic short side calculation.
+
+        Examples:
+            >>> transform = ClassifyLetterBox(size=224)
+            >>> img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
+            >>> result = transform(img)
+            >>> print(result.shape)
+            (224, 224, 3)
         """
         super().__init__()
         self.h, self.w = (size, size) if isinstance(size, int) else size
@@ -1175,13 +2561,24 @@ class ClassifyLetterBox:
 
     def __call__(self, im):
         """
-        Resizes the image and pads it with a letterbox method.
+        Resizes and pads an image using the letterbox method.
+
+        This method resizes the input image to fit within the specified dimensions while maintaining its aspect ratio,
+        then pads the resized image to match the target size.
 
         Args:
-            im (numpy.ndarray): The input image as a numpy array of shape HWC.
+            im (numpy.ndarray): Input image as a numpy array with shape (H, W, C).
 
         Returns:
-            (numpy.ndarray): The letterboxed and resized image as a numpy array.
+            (numpy.ndarray): Resized and padded image as a numpy array with shape (hs, ws, 3), where hs and ws are
+                the target height and width respectively.
+
+        Examples:
+            >>> letterbox = ClassifyLetterBox(size=(640, 640))
+            >>> image = np.random.randint(0, 255, (720, 1280, 3), dtype=np.uint8)
+            >>> resized_image = letterbox(image)
+            >>> print(resized_image.shape)
+            (640, 640, 3)
         """
         imh, imw = im.shape[:2]
         r = min(self.h / imh, self.w / imw)  # ratio of new/old dimensions
@@ -1199,25 +2596,73 @@ class ClassifyLetterBox:
 
 # NOTE: keep this class for backward compatibility
 class CenterCrop:
-    """YOLOv8 CenterCrop class for image preprocessing, designed to be part of a transformation pipeline, e.g.,
-    T.Compose([CenterCrop(size), ToTensor()]).
+    """
+    Applies center cropping to images for classification tasks.
+
+    This class performs center cropping on input images, resizing them to a specified size while maintaining the aspect
+    ratio. It is designed to be part of a transformation pipeline, e.g., T.Compose([CenterCrop(size), ToTensor()]).
+
+    Attributes:
+        h (int): Target height of the cropped image.
+        w (int): Target width of the cropped image.
+
+    Methods:
+        __call__: Applies the center crop transformation to an input image.
+
+    Examples:
+        >>> transform = CenterCrop(640)
+        >>> image = np.random.randint(0, 255, (1080, 1920, 3), dtype=np.uint8)
+        >>> cropped_image = transform(image)
+        >>> print(cropped_image.shape)
+        (640, 640, 3)
     """
 
     def __init__(self, size=640):
-        """Converts an image from numpy array to PyTorch tensor."""
+        """
+        Initializes the CenterCrop object for image preprocessing.
+
+        This class is designed to be part of a transformation pipeline, e.g., T.Compose([CenterCrop(size), ToTensor()]).
+        It performs a center crop on input images to a specified size.
+
+        Args:
+            size (int | Tuple[int, int]): The desired output size of the crop. If size is an int, a square crop
+                (size, size) is made. If size is a sequence like (h, w), it is used as the output size.
+
+        Returns:
+            (None): This method initializes the object and does not return anything.
+
+        Examples:
+            >>> transform = CenterCrop(224)
+            >>> img = np.random.rand(300, 300, 3)
+            >>> cropped_img = transform(img)
+            >>> print(cropped_img.shape)
+            (224, 224, 3)
+        """
         super().__init__()
         self.h, self.w = (size, size) if isinstance(size, int) else size
 
     def __call__(self, im):
         """
-        Resizes and crops the center of the image using a letterbox method.
+        Applies center cropping to an input image.
+
+        This method resizes and crops the center of the image using a letterbox method. It maintains the aspect
+        ratio of the original image while fitting it into the specified dimensions.
 
         Args:
-            im (numpy.ndarray): The input image as a numpy array of shape HWC.
+            im (numpy.ndarray | PIL.Image.Image): The input image as a numpy array of shape (H, W, C) or a
+                PIL Image object.
 
         Returns:
-            (numpy.ndarray): The center-cropped and resized image as a numpy array.
+            (numpy.ndarray): The center-cropped and resized image as a numpy array of shape (self.h, self.w, C).
+
+        Examples:
+            >>> transform = CenterCrop(size=224)
+            >>> image = np.random.randint(0, 255, (640, 480, 3), dtype=np.uint8)
+            >>> cropped_image = transform(image)
+            >>> assert cropped_image.shape == (224, 224, 3)
         """
+        if isinstance(im, Image.Image):  # convert from PIL to numpy array if required
+            im = np.asarray(im)
         imh, imw = im.shape[:2]
         m = min(imh, imw)  # min dimension
         top, left = (imh - m) // 2, (imw - m) // 2
@@ -1226,22 +2671,71 @@ class CenterCrop:
 
 # NOTE: keep this class for backward compatibility
 class ToTensor:
-    """YOLOv8 ToTensor class for image preprocessing, i.e., T.Compose([LetterBox(size), ToTensor()])."""
+    """
+    Converts an image from a numpy array to a PyTorch tensor.
+
+    This class is designed to be part of a transformation pipeline, e.g., T.Compose([LetterBox(size), ToTensor()]).
+
+    Attributes:
+        half (bool): If True, converts the image to half precision (float16).
+
+    Methods:
+        __call__: Applies the tensor conversion to an input image.
+
+    Examples:
+        >>> transform = ToTensor(half=True)
+        >>> img = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
+        >>> tensor_img = transform(img)
+        >>> print(tensor_img.shape, tensor_img.dtype)
+        torch.Size([3, 640, 640]) torch.float16
+
+    Notes:
+        The input image is expected to be in BGR format with shape (H, W, C).
+        The output tensor will be in RGB format with shape (C, H, W), normalized to [0, 1].
+    """
 
     def __init__(self, half=False):
-        """Initialize YOLOv8 ToTensor object with optional half-precision support."""
+        """
+        Initializes the ToTensor object for converting images to PyTorch tensors.
+
+        This class is designed to be used as part of a transformation pipeline for image preprocessing in the
+        Ultralytics YOLO framework. It converts numpy arrays or PIL Images to PyTorch tensors, with an option
+        for half-precision (float16) conversion.
+
+        Args:
+            half (bool): If True, converts the tensor to half precision (float16). Default is False.
+
+        Examples:
+            >>> transform = ToTensor(half=True)
+            >>> img = np.random.rand(640, 640, 3)
+            >>> tensor_img = transform(img)
+            >>> print(tensor_img.dtype)
+            torch.float16
+        """
         super().__init__()
         self.half = half
 
     def __call__(self, im):
         """
-        Transforms an image from a numpy array to a PyTorch tensor, applying optional half-precision and normalization.
+        Transforms an image from a numpy array to a PyTorch tensor.
+
+        This method converts the input image from a numpy array to a PyTorch tensor, applying optional
+        half-precision conversion and normalization. The image is transposed from HWC to CHW format and
+        the color channels are reversed from BGR to RGB.
 
         Args:
             im (numpy.ndarray): Input image as a numpy array with shape (H, W, C) in BGR order.
 
         Returns:
-            (torch.Tensor): The transformed image as a PyTorch tensor in float32 or float16, normalized to [0, 1].
+            (torch.Tensor): The transformed image as a PyTorch tensor in float32 or float16, normalized
+                to [0, 1] with shape (C, H, W) in RGB order.
+
+        Examples:
+            >>> transform = ToTensor(half=True)
+            >>> img = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
+            >>> tensor_img = transform(img)
+            >>> print(tensor_img.shape, tensor_img.dtype)
+            torch.Size([3, 640, 640]) torch.float16
         """
         im = np.ascontiguousarray(im.transpose((2, 0, 1))[::-1])  # HWC to CHW -> BGR to RGB -> contiguous
         im = torch.from_numpy(im)  # to torch