dgenerate-ultralytics-headless 8.3.196__py3-none-any.whl → 8.3.248__py3-none-any.whl

ultralytics/data/augment.py

@@ -26,11 +26,10 @@ DEFAULT_STD = (1.0, 1.0, 1.0)
 
 
 class BaseTransform:
-    """
-    Base class for image transformations in the Ultralytics library.
+    """Base class for image transformations in the Ultralytics library.
 
-    This class serves as a foundation for implementing various image processing operations, 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:
         apply_image: Apply image transformations to labels.
@@ -45,27 +44,22 @@ class BaseTransform:
     """
 
     def __init__(self) -> None:
-        """
-        Initialize 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.
+        """Initialize the BaseTransform object.
 
-        Examples:
-            >>> transform = BaseTransform()
+        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.
         """
         pass
 
     def apply_image(self, labels):
-        """
-        Apply image transformations to labels.
+        """Apply 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.
+            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.
@@ -80,8 +74,7 @@ class BaseTransform:
         pass
 
     def apply_instances(self, labels):
-        """
-        Apply transformations to object instances in labels.
+        """Apply 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
@@ -101,8 +94,7 @@ class BaseTransform:
         pass
 
     def apply_semantic(self, labels):
-        """
-        Apply semantic segmentation transformations to an image.
+        """Apply 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.
@@ -121,16 +113,15 @@ class BaseTransform:
         pass
 
     def __call__(self, labels):
-        """
-        Apply all label transformations to an image, instances, and semantic masks.
+        """Apply 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.
+        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.
+            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.
@@ -146,8 +137,7 @@ class BaseTransform:
 
 
 class Compose:
-    """
-    A 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.
@@ -169,28 +159,21 @@ class Compose:
     """
 
     def __init__(self, transforms):
-        """
-        Initialize the Compose object with a list of transforms.
+        """Initialize 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):
-        """
-        Apply a series of transformations to input data.
+        """Apply a series of transformations to input data.
 
         This method sequentially applies each transformation in the Compose object's 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.
+            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.
@@ -205,8 +188,7 @@ class Compose:
         return data
 
     def append(self, transform):
-        """
-        Append a new transform to the existing list of transforms.
+        """Append a new transform to the existing list of transforms.
 
         Args:
             transform (BaseTransform): The transformation to be added to the composition.
@@ -218,8 +200,7 @@ class Compose:
         self.transforms.append(transform)
 
     def insert(self, index, transform):
-        """
-        Insert a new transform at a specified index in the existing list of transforms.
+        """Insert 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.
@@ -234,8 +215,7 @@ class Compose:
         self.transforms.insert(index, transform)
 
     def __getitem__(self, index: list | int) -> Compose:
-        """
-        Retrieve a specific transform or a set of transforms using indexing.
+        """Retrieve 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.
@@ -256,8 +236,7 @@ class Compose:
         return Compose([self.transforms[i] for i in index]) if isinstance(index, list) else self.transforms[index]
 
     def __setitem__(self, index: list | int, value: list | int) -> None:
-        """
-        Set one or more transforms in the composition using indexing.
+        """Set one or more transforms in the composition using indexing.
 
         Args:
             index (int | list[int]): Index or list of indices to set transforms at.
@@ -283,8 +262,7 @@ class Compose:
             self.transforms[i] = v
 
     def tolist(self):
-        """
-        Convert the list of transforms to a standard Python list.
+        """Convert the list of transforms to a standard Python list.
 
         Returns:
             (list): A list containing all the transform objects in the Compose instance.
@@ -299,8 +277,7 @@ class Compose:
         return self.transforms
 
     def __repr__(self):
-        """
-        Return a string representation of the Compose object.
+        """Return a string representation of the Compose object.
 
         Returns:
             (str): A string representation of the Compose object, including the list of transforms.
@@ -318,11 +295,10 @@ class Compose:
 
 
 class BaseMixTransform:
-    """
-    Base class for mix transformations like Cutmix, MixUp and Mosaic.
+    """Base class for mix transformations like Cutmix, 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 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.
 
     Attributes:
         dataset (Any): The dataset object containing images and labels.
@@ -349,8 +325,7 @@ class BaseMixTransform:
     """
 
     def __init__(self, dataset, pre_transform=None, p=0.0) -> None:
-        """
-        Initialize the BaseMixTransform object for mix transformations like CutMix, MixUp and Mosaic.
+        """Initialize the BaseMixTransform object for mix transformations like CutMix, MixUp and Mosaic.
 
         This class serves as a base for implementing mix transformations in image processing pipelines.
 
@@ -358,22 +333,16 @@ class BaseMixTransform:
             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: dict[str, Any]) -> dict[str, Any]:
-        """
-        Apply pre-processing transforms and cutmix/mixup/mosaic transforms to labels data.
+        """Apply pre-processing transforms and cutmix/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.
+        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[str, Any]): A dictionary containing label data for an image.
@@ -409,8 +378,7 @@ class BaseMixTransform:
         return labels
 
     def _mix_transform(self, labels: dict[str, Any]):
-        """
-        Apply CutMix, MixUp or Mosaic augmentation to the label dictionary.
+        """Apply CutMix, MixUp or Mosaic augmentation to the label dictionary.
 
         This method should be implemented by subclasses to perform specific mix transformations like CutMix, MixUp or
         Mosaic. It modifies the input label dictionary in-place with the augmented data.
@@ -430,8 +398,7 @@ class BaseMixTransform:
         raise NotImplementedError
 
     def get_indexes(self):
-        """
-        Get a list of shuffled indexes for mosaic augmentation.
+        """Get a list of shuffled indexes for mosaic augmentation.
 
         Returns:
             (list[int]): A list of shuffled indexes from the dataset.
@@ -445,15 +412,14 @@ class BaseMixTransform:
 
     @staticmethod
     def _update_label_text(labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Update label text and class IDs for mixed labels in image augmentation.
+        """Update 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.
+        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[str, Any]): A dictionary containing label information, including 'texts' and 'cls' fields,
-                and optionally a 'mix_labels' field with additional label dictionaries.
+            labels (dict[str, Any]): A dictionary containing label information, including 'texts' and 'cls' fields, and
+                optionally a 'mix_labels' field with additional label dictionaries.
 
         Returns:
             (dict[str, Any]): The updated labels dictionary with unified text labels and updated class IDs.
@@ -477,7 +443,7 @@ class BaseMixTransform:
         if "texts" not in labels:
             return labels
 
-        mix_texts = sum([labels["texts"]] + [x["texts"] for x in labels["mix_labels"]], [])
+        mix_texts = [*labels["texts"], *(item for x in labels["mix_labels"] for item in x["texts"])]
         mix_texts = list({tuple(x) for x in mix_texts})
         text2id = {text: i for i, text in enumerate(mix_texts)}
 
@@ -490,11 +456,10 @@ class BaseMixTransform:
 
 
 class Mosaic(BaseMixTransform):
-    """
-    Mosaic augmentation for image datasets.
+    """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.
+    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.
@@ -520,22 +485,16 @@ class Mosaic(BaseMixTransform):
     """
 
     def __init__(self, dataset, imgsz: int = 640, p: float = 1.0, n: int = 4):
-        """
-        Initialize the Mosaic augmentation object.
+        """Initialize 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.
+        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."
@@ -546,15 +505,14 @@ class Mosaic(BaseMixTransform):
         self.buffer_enabled = self.dataset.cache != "ram"
 
     def get_indexes(self):
-        """
-        Return a list of random indexes from the dataset for mosaic augmentation.
+        """Return 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.
+        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.
 
         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).
+            (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)
@@ -567,12 +525,11 @@ class Mosaic(BaseMixTransform):
             return [random.randint(0, len(self.dataset) - 1) for _ in range(self.n - 1)]
 
     def _mix_transform(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Apply mosaic augmentation to the input image and labels.
+        """Apply 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.
+        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[str, Any]): A dictionary containing image data and annotations. Expected keys include:
@@ -596,16 +553,15 @@ class Mosaic(BaseMixTransform):
         )  # This code is modified for mosaic3 method.
 
     def _mosaic3(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Create a 1x3 image mosaic by combining three images.
+        """Create 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.
+        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[str, Any]): 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.
+                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[str, Any]): A dictionary with the mosaic image and updated labels. Keys include:
@@ -655,19 +611,19 @@ class Mosaic(BaseMixTransform):
         return final_labels
 
     def _mosaic4(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Create a 2x2 image mosaic from four input images.
+        """Create 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.
+        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[str, Any]): 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.
+            labels (dict[str, Any]): 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[str, Any]): 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.
+            (dict[str, Any]): 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)
@@ -713,22 +669,22 @@ class Mosaic(BaseMixTransform):
         return final_labels
 
     def _mosaic9(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Create a 3x3 image mosaic from the input image and eight additional images.
+        """Create 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.
+        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[str, Any]): A dictionary containing the input image and its associated labels. It should have
-                the following keys:
+            the following keys:
                 - 'img' (np.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.
+            eight images, each with the same structure as the input labels.
 
         Returns:
-            (dict[str, Any]): A dictionary containing the mosaic image and updated labels. It includes the following keys:
+            (dict[str, Any]): A dictionary containing the mosaic image and updated labels. It includes the following
+            keys:
                 - 'img' (np.ndarray): The final mosaic image.
                 - Other keys from the input labels, updated to reflect the new mosaic arrangement.
 
@@ -786,8 +742,7 @@ class Mosaic(BaseMixTransform):
 
     @staticmethod
     def _update_labels(labels, padw: int, padh: int) -> dict[str, Any]:
-        """
-        Update label coordinates with padding values.
+        """Update 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.
@@ -812,11 +767,10 @@ class Mosaic(BaseMixTransform):
         return labels
 
     def _cat_labels(self, mosaic_labels: list[dict[str, Any]]) -> dict[str, Any]:
-        """
-        Concatenate and process labels for mosaic augmentation.
+        """Concatenate and process 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.
+        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[str, Any]]): A list of label dictionaries for each image in the mosaic.
@@ -864,8 +818,7 @@ class Mosaic(BaseMixTransform):
 
 
 class MixUp(BaseMixTransform):
-    """
-    Apply MixUp augmentation to image datasets.
+    """Apply 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.
@@ -886,30 +839,23 @@ class MixUp(BaseMixTransform):
     """
 
     def __init__(self, dataset, pre_transform=None, p: float = 0.0) -> None:
-        """
-        Initialize the MixUp augmentation object.
+        """Initialize 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.
+        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 _mix_transform(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Apply MixUp augmentation to the input labels.
+        """Apply 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).
+        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[str, Any]): A dictionary containing the original image and label information.
@@ -930,11 +876,10 @@ class MixUp(BaseMixTransform):
 
 
 class CutMix(BaseMixTransform):
-    """
-    Apply CutMix augmentation to image datasets as described in the paper https://arxiv.org/abs/1905.04899.
+    """Apply CutMix augmentation to image datasets as described in the paper https://arxiv.org/abs/1905.04899.
 
-    CutMix combines two images by replacing a random rectangular region of one image with the corresponding region from another image,
-    and adjusts the labels proportionally to the area of the mixed region.
+    CutMix combines two images by replacing a random rectangular region of one image with the corresponding region from
+    another image, and adjusts the labels proportionally to the area of the mixed region.
 
     Attributes:
         dataset (Any): The dataset to which CutMix augmentation will be applied.
@@ -955,8 +900,7 @@ class CutMix(BaseMixTransform):
     """
 
     def __init__(self, dataset, pre_transform=None, p: float = 0.0, beta: float = 1.0, num_areas: int = 3) -> None:
-        """
-        Initialize the CutMix augmentation object.
+        """Initialize the CutMix augmentation object.
 
         Args:
             dataset (Any): The dataset to which CutMix augmentation will be applied.
@@ -970,8 +914,7 @@ class CutMix(BaseMixTransform):
         self.num_areas = num_areas
 
     def _rand_bbox(self, width: int, height: int) -> tuple[int, int, int, int]:
-        """
-        Generate random bounding box coordinates for the cut region.
+        """Generate random bounding box coordinates for the cut region.
 
         Args:
             width (int): Width of the image.
@@ -1000,8 +943,7 @@ class CutMix(BaseMixTransform):
         return x1, y1, x2, y2
 
     def _mix_transform(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Apply CutMix augmentation to the input labels.
+        """Apply CutMix augmentation to the input labels.
 
         Args:
             labels (dict[str, Any]): A dictionary containing the original image and label information.
@@ -1048,12 +990,11 @@ class CutMix(BaseMixTransform):
 
 
 class RandomPerspective:
-    """
-    Implement random perspective and affine transformations on images and corresponding annotations.
+    """Implement 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.
+    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): Maximum absolute degree range for random rotations.
@@ -1091,8 +1032,7 @@ class RandomPerspective:
         border: tuple[int, int] = (0, 0),
         pre_transform=None,
     ):
-        """
-        Initialize RandomPerspective object with transformation parameters.
+        """Initialize 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.
@@ -1106,10 +1046,6 @@ class RandomPerspective:
             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
@@ -1120,12 +1056,11 @@ class RandomPerspective:
         self.pre_transform = pre_transform
 
     def affine_transform(self, img: np.ndarray, border: tuple[int, int]) -> tuple[np.ndarray, np.ndarray, float]:
-        """
-        Apply a sequence of affine transformations centered around the image center.
+        """Apply 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.
+        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 (np.ndarray): Input image to be transformed.
@@ -1184,15 +1119,14 @@ class RandomPerspective:
         return img, M, s
 
     def apply_bboxes(self, bboxes: np.ndarray, M: np.ndarray) -> np.ndarray:
-        """
-        Apply affine transformation to bounding boxes.
+        """Apply affine transformation to bounding boxes.
 
-        This function applies an affine transformation to a set of bounding boxes using the provided
-        transformation matrix.
+        This function applies an affine transformation to a set of bounding boxes using the provided transformation
+        matrix.
 
         Args:
-            bboxes (np.ndarray): Bounding boxes in xyxy format with shape (N, 4), where N is the number
-                of bounding boxes.
+            bboxes (np.ndarray): Bounding boxes in xyxy format with shape (N, 4), where N is the number of bounding
+                boxes.
             M (np.ndarray): Affine transformation matrix with shape (3, 3).
 
         Returns:
@@ -1218,11 +1152,10 @@ class RandomPerspective:
         return np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1)), dtype=bboxes.dtype).reshape(4, n).T
 
     def apply_segments(self, segments: np.ndarray, M: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-        """
-        Apply affine transformations to segments and generate new bounding boxes.
+        """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.
+        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 (np.ndarray): Input segments with shape (N, M, 2), where N is the number of segments and M is the
@@ -1254,16 +1187,15 @@ class RandomPerspective:
         return bboxes, segments
 
     def apply_keypoints(self, keypoints: np.ndarray, M: np.ndarray) -> np.ndarray:
-        """
-        Apply affine transformation to keypoints.
+        """Apply 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 (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).
+            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:
@@ -1288,21 +1220,14 @@ class RandomPerspective:
         return np.concatenate([xy, visible], axis=-1).reshape(n, nkpt, 3)
 
     def __call__(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Apply random perspective and affine transformations to an image and its associated labels.
+        """Apply 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.
+        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[str, Any]): A dictionary containing image data and annotations.
-                Must include:
-                    'img' (np.ndarray): The input image.
-                    'cls' (np.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[str, Any]): Transformed labels dictionary containing:
@@ -1321,6 +1246,14 @@ class RandomPerspective:
             ... }
             >>> result = transform(labels)
             >>> assert result["img"].shape[:2] == result["resized_shape"]
+
+        Notes:
+            'labels' arg must include:
+                - 'img' (np.ndarray): The input image.
+                - 'cls' (np.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.
         """
         if self.pre_transform and "mosaic_border" not in labels:
             labels = self.pre_transform(labels)
@@ -1374,29 +1307,27 @@ class RandomPerspective:
         area_thr: float = 0.1,
         eps: float = 1e-16,
     ) -> np.ndarray:
-        """
-        Compute candidate boxes for further processing based on size and aspect ratio criteria.
+        """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.
+        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 (np.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 (np.ndarray): Augmented boxes after transformation, shape (4, N). Format is
-                [x1, y1, x2, y2] in absolute coordinates.
-            wh_thr (int): Width and height threshold in pixels. Boxes smaller than this in either
-                dimension are rejected.
-            ar_thr (int): 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.
+            box1 (np.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 (np.ndarray): Augmented boxes after transformation, shape (4, N). Format is [x1, y1, x2, y2] in
+                absolute coordinates.
+            wh_thr (int): Width and height threshold in pixels. Boxes smaller than this in either dimension are
+                rejected.
+            ar_thr (int): 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:
-            (np.ndarray): Boolean array of shape (n) indicating which boxes are candidates.
-                True values correspond to boxes that meet all criteria.
+            (np.ndarray): Boolean array of shape (n) indicating which boxes are candidates. True values correspond to
+                boxes that meet all criteria.
 
         Examples:
             >>> random_perspective = RandomPerspective()
@@ -1413,8 +1344,7 @@ class RandomPerspective:
 
 
 class RandomHSV:
-    """
-    Randomly adjust the Hue, Saturation, and Value (HSV) channels of an image.
+    """Randomly adjust 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.
 
@@ -1437,8 +1367,7 @@ class RandomHSV:
     """
 
     def __init__(self, hgain: float = 0.5, sgain: float = 0.5, vgain: float = 0.5) -> None:
-        """
-        Initialize the RandomHSV object for random HSV (Hue, Saturation, Value) augmentation.
+        """Initialize 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.
 
@@ -1446,25 +1375,20 @@ class RandomHSV:
             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
         self.vgain = vgain
 
     def __call__(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Apply random HSV augmentation to an image within predefined limits.
+        """Apply random HSV augmentation to an image within predefined limits.
 
-        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.
+        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[str, Any]): A dictionary containing image data and metadata. Must include an 'img' key with
-                the image as a numpy array.
+            labels (dict[str, Any]): A dictionary containing image data and metadata. Must include an 'img' key with the
+                image as a numpy array.
 
         Returns:
             (dict[str, Any]): A dictionary containing the mixed image and adjusted labels.
@@ -1496,11 +1420,10 @@ class RandomHSV:
 
 
 class RandomFlip:
-    """
-    Apply a random horizontal or vertical flip to an image with a given probability.
+    """Apply a random horizontal or vertical flip to an image with a given probability.
 
-    This class performs random image flipping and updates corresponding instance annotations such as
-    bounding boxes and keypoints.
+    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.
@@ -1517,12 +1440,11 @@ class RandomFlip:
         >>> flipped_instances = result["instances"]
     """
 
-    def __init__(self, p: float = 0.5, direction: str = "horizontal", flip_idx: list[int] = None) -> None:
-        """
-        Initialize the RandomFlip class with probability and direction.
+    def __init__(self, p: float = 0.5, direction: str = "horizontal", flip_idx: list[int] | None = None) -> None:
+        """Initialize 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.
+        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): The probability of applying the flip. Must be between 0 and 1.
@@ -1531,10 +1453,6 @@ class RandomFlip:
 
         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, f"The probability should be in range [0, 1], but got {p}."
@@ -1544,23 +1462,21 @@ class RandomFlip:
         self.flip_idx = flip_idx
 
     def __call__(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Apply random flip to an image and update any instances like bounding boxes or keypoints accordingly.
+        """Apply random flip to an image and update 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.
+        probability and direction. It also updates the corresponding instances (bounding boxes, keypoints) to match the
+        flipped image.
 
         Args:
             labels (dict[str, Any]): A dictionary containing the following keys:
-                'img' (np.ndarray): The image to be flipped.
-                'instances' (ultralytics.utils.instance.Instances): An object containing bounding boxes and
-                    optionally keypoints.
+                - 'img' (np.ndarray): The image to be flipped.
+                - 'instances' (ultralytics.utils.instance.Instances): Object containing boxes and optionally keypoints.
 
         Returns:
             (dict[str, Any]): The same dictionary with the flipped image and updated instances:
-                'img' (np.ndarray): The flipped image.
-                'instances' (ultralytics.utils.instance.Instances): Updated instances matching the flipped image.
+                - 'img' (np.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(...)}
@@ -1591,11 +1507,10 @@ 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.
+    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.
@@ -1626,8 +1541,7 @@ class LetterBox:
         padding_value: int = 114,
         interpolation: int = cv2.INTER_LINEAR,
     ):
-        """
-        Initialize LetterBox object for resizing and padding images.
+        """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.
@@ -1641,19 +1555,6 @@ class LetterBox:
             stride (int): Stride of the model (e.g., 32 for YOLOv5).
             padding_value (int): Value for padding the image. Default is 114.
             interpolation (int): Interpolation method for resizing. Default is cv2.INTER_LINEAR.
-
-        Attributes:
-            new_shape (tuple[int, int]): Target size for the resized image.
-            auto (bool): Flag for using minimum rectangle resizing.
-            scale_fill (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.
-            padding_value (int): Value used for padding the image.
-            interpolation (int): Interpolation method used for resizing.
-
-        Examples:
-            >>> letterbox = LetterBox(new_shape=(640, 640), auto=False, scale_fill=False, scaleup=True, stride=32)
-            >>> resized_img = letterbox(original_img)
         """
         self.new_shape = new_shape
         self.auto = auto
@@ -1664,21 +1565,21 @@ class LetterBox:
         self.padding_value = padding_value
         self.interpolation = interpolation
 
-    def __call__(self, labels: dict[str, Any] = None, image: np.ndarray = None) -> dict[str, Any] | np.ndarray:
-        """
-        Resize and pad an image for object detection, instance segmentation, or pose estimation tasks.
+    def __call__(self, labels: dict[str, Any] | None = None, image: np.ndarray = None) -> dict[str, Any] | np.ndarray:
+        """Resize and pad 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[str, Any] | None): A dictionary containing image data and associated labels, or empty dict if None.
+            labels (dict[str, Any] | 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[str, Any] | nd.ndarray): If 'labels' is provided, returns an updated dictionary with the resized and padded image,
-                updated labels, and additional metadata. If 'labels' is empty, returns the resized
-                and padded image.
+            (dict[str, Any] | np.ndarray): If 'labels' is provided, returns an updated dictionary with the resized and
+                padded image, updated labels, and additional metadata. If 'labels' is empty, returns the resized and
+                padded image only.
 
         Examples:
             >>> letterbox = LetterBox(new_shape=(640, 640))
@@ -1701,7 +1602,7 @@ class LetterBox:
 
         # Compute padding
         ratio = r, r  # width, height ratios
-        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
+        new_unpad = round(shape[1] * r), round(shape[0] * r)
         dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
         if self.auto:  # minimum rectangle
             dw, dh = np.mod(dw, self.stride), np.mod(dh, self.stride)  # wh padding
@@ -1719,8 +1620,8 @@ class LetterBox:
             if img.ndim == 2:
                 img = img[..., None]
 
-        top, bottom = int(round(dh - 0.1)) if self.center else 0, int(round(dh + 0.1))
-        left, right = int(round(dw - 0.1)) if self.center else 0, int(round(dw + 0.1))
+        top, bottom = round(dh - 0.1) if self.center else 0, round(dh + 0.1)
+        left, right = round(dw - 0.1) if self.center else 0, round(dw + 0.1)
         h, w, c = img.shape
         if c == 3:
             img = cv2.copyMakeBorder(
@@ -1744,11 +1645,10 @@ class LetterBox:
 
     @staticmethod
     def _update_labels(labels: dict[str, Any], ratio: tuple[float, float], padw: float, padh: float) -> dict[str, Any]:
-        """
-        Update labels after applying letterboxing to an image.
+        """Update 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.
+        This method modifies the bounding box coordinates of instances in the labels to account for resizing and padding
+        applied during letterboxing.
 
         Args:
             labels (dict[str, Any]): A dictionary containing image labels and instances.
@@ -1774,8 +1674,7 @@ class LetterBox:
 
 
 class CopyPaste(BaseMixTransform):
-    """
-    CopyPaste class for applying Copy-Paste augmentation to image datasets.
+    """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
@@ -1798,7 +1697,7 @@ class CopyPaste(BaseMixTransform):
     """
 
     def __init__(self, dataset=None, pre_transform=None, p: float = 0.5, mode: str = "flip") -> None:
-        """Initialize CopyPaste object with dataset, pre_transform, and probability of applying MixUp."""
+        """Initialize CopyPaste object with dataset, pre_transform, and probability of applying CopyPaste."""
         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
@@ -1874,8 +1773,7 @@ class CopyPaste(BaseMixTransform):
 
 
 class Albumentations:
-    """
-    Albumentations transformations for image augmentation.
+    """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
@@ -1894,14 +1792,13 @@ class Albumentations:
         >>> 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.
+        - 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.
     """
 
-    def __init__(self, p: float = 1.0) -> None:
-        """
-        Initialize the Albumentations transform object for YOLO bbox formatted parameters.
+    def __init__(self, p: float = 1.0, transforms: list | None = None) -> None:
+        """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
@@ -1909,26 +1806,11 @@ class Albumentations:
 
         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.
+            transforms (list, optional): List of custom Albumentations transforms. If None, uses default transforms.
 
         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
@@ -1986,16 +1868,20 @@ class Albumentations:
                 "XYMasking",
             }  # from https://albumentations.ai/docs/getting_started/transforms_and_targets/#spatial-level-transforms
 
-            # Transforms
-            T = [
-                A.Blur(p=0.01),
-                A.MedianBlur(p=0.01),
-                A.ToGray(p=0.01),
-                A.CLAHE(p=0.01),
-                A.RandomBrightnessContrast(p=0.0),
-                A.RandomGamma(p=0.0),
-                A.ImageCompression(quality_range=(75, 100), p=0.0),
-            ]
+            # Transforms, use custom transforms if provided, otherwise use defaults
+            T = (
+                [
+                    A.Blur(p=0.01),
+                    A.MedianBlur(p=0.01),
+                    A.ToGray(p=0.01),
+                    A.CLAHE(p=0.01),
+                    A.RandomBrightnessContrast(p=0.0),
+                    A.RandomGamma(p=0.0),
+                    A.ImageCompression(quality_range=(75, 100), p=0.0),
+                ]
+                if transforms is None
+                else transforms
+            )
 
             # Compose transforms
             self.contains_spatial = any(transform.__class__.__name__ in spatial_transforms for transform in T)
@@ -2014,8 +1900,7 @@ class Albumentations:
             LOGGER.info(f"{prefix}{e}")
 
     def __call__(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Apply Albumentations transformations to input labels.
+        """Apply 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.
@@ -2061,7 +1946,7 @@ class Albumentations:
                 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"])
+                    labels["cls"] = np.array(new["class_labels"]).reshape(-1, 1)
                     bboxes = np.array(new["bboxes"], dtype=np.float32)
                 labels["instances"].update(bboxes=bboxes)
         else:
@@ -2071,8 +1956,7 @@ class Albumentations:
 
 
 class Format:
-    """
-    A class for formatting image annotations for object detection, instance segmentation, and pose estimation tasks.
+    """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.
 
@@ -2112,8 +1996,7 @@ class Format:
         batch_idx: bool = True,
         bgr: float = 0.0,
     ):
-        """
-        Initialize the Format class with given parameters for image and instance annotation formatting.
+        """Initialize 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`.
@@ -2128,22 +2011,6 @@ class Format:
             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
@@ -2156,8 +2023,7 @@ class Format:
         self.bgr = bgr
 
     def __call__(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Format image annotations for object detection, instance segmentation, and pose estimation tasks.
+        """Format 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
@@ -2225,8 +2091,7 @@ class Format:
         return labels
 
     def _format_img(self, img: np.ndarray) -> torch.Tensor:
-        """
-        Format an image for YOLO from a Numpy array to a PyTorch tensor.
+        """Format 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).
@@ -2249,7 +2114,7 @@ class Format:
             torch.Size([3, 100, 100])
         """
         if len(img.shape) < 3:
-            img = np.expand_dims(img, -1)
+            img = img[..., None]
         img = img.transpose(2, 0, 1)
         img = np.ascontiguousarray(img[::-1] if random.uniform(0, 1) > self.bgr and img.shape[0] == 3 else img)
         img = torch.from_numpy(img)
@@ -2258,8 +2123,7 @@ class Format:
     def _format_segments(
         self, instances: Instances, cls: np.ndarray, w: int, h: int
     ) -> tuple[np.ndarray, Instances, np.ndarray]:
-        """
-        Convert polygon segments to bitmap masks.
+        """Convert polygon segments to bitmap masks.
 
         Args:
             instances (Instances): Object containing segment information.
@@ -2293,17 +2157,16 @@ class LoadVisualPrompt:
     """Create visual prompts from bounding boxes or masks for model input."""
 
     def __init__(self, scale_factor: float = 1 / 8) -> None:
-        """
-        Initialize the LoadVisualPrompt with a scale factor.
+        """Initialize the LoadVisualPrompt with a scale factor.
 
         Args:
             scale_factor (float): Factor to scale the input image dimensions.
         """
         self.scale_factor = scale_factor
 
-    def make_mask(self, boxes: torch.Tensor, h: int, w: int) -> torch.Tensor:
-        """
-        Create binary masks from bounding boxes.
+    @staticmethod
+    def make_mask(boxes: torch.Tensor, h: int, w: int) -> torch.Tensor:
+        """Create binary masks from bounding boxes.
 
         Args:
             boxes (torch.Tensor): Bounding boxes in xyxy format, shape: (N, 4).
@@ -2320,8 +2183,7 @@ class LoadVisualPrompt:
         return (r >= x1) * (r < x2) * (c >= y1) * (c < y2)
 
     def __call__(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Process labels to create visual prompts.
+        """Process labels to create visual prompts.
 
         Args:
             labels (dict[str, Any]): Dictionary containing image data and annotations.
@@ -2347,8 +2209,7 @@ class LoadVisualPrompt:
         bboxes: np.ndarray | torch.Tensor = None,
         masks: np.ndarray | torch.Tensor = None,
     ) -> torch.Tensor:
-        """
-        Generate visual masks based on bounding boxes or masks.
+        """Generate visual masks based on bounding boxes or masks.
 
         Args:
             category (int | np.ndarray | torch.Tensor): The category labels for the objects.
@@ -2382,19 +2243,18 @@ class LoadVisualPrompt:
         #     assert len(cls_unique) == cls_unique[-1] + 1, (
         #         f"Expected a continuous range of class indices, but got {cls_unique}"
         #     )
-        visuals = torch.zeros(len(cls_unique), *masksz)
+        visuals = torch.zeros(cls_unique.shape[0], *masksz)
         for idx, mask in zip(inverse_indices, masks):
             visuals[idx] = torch.logical_or(visuals[idx], mask)
         return visuals
 
 
 class RandomLoadText:
-    """
-    Randomly sample positive and negative texts and update class indices accordingly.
+    """Randomly sample positive and negative texts and update 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.
+    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.
@@ -2422,38 +2282,20 @@ class RandomLoadText:
         padding: bool = False,
         padding_value: list[str] = [""],
     ) -> None:
-        """
-        Initialize the RandomLoadText class for randomly sampling positive and negative texts.
+        """Initialize 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.
+        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. 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.
+            prompt_format (str): Format string for the prompt. 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.
             max_samples (int): The maximum number of different text samples in one image.
-            padding (bool): Whether to pad texts to max_samples. If True, the number of texts will always
-                be equal to max_samples.
+            padding (bool): Whether to pad texts to max_samples. If True, the number of texts will always be equal to
+                max_samples.
             padding_value (str): The padding text to use when padding is True.
-
-        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
@@ -2462,15 +2304,15 @@ class RandomLoadText:
         self.padding_value = padding_value
 
     def __call__(self, labels: dict[str, Any]) -> dict[str, Any]:
-        """
-        Randomly sample positive and negative texts and update class indices accordingly.
+        """Randomly sample positive and negative texts and update 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.
+        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[str, Any]): A dictionary containing image labels and metadata. Must include 'texts' and 'cls' keys.
+            labels (dict[str, Any]): A dictionary containing image labels and metadata. Must include 'texts' and 'cls'
+                keys.
 
         Returns:
             (dict[str, Any]): Updated labels dictionary with new 'cls' and 'texts' entries.
@@ -2528,16 +2370,16 @@ class RandomLoadText:
 
 
 def v8_transforms(dataset, imgsz: int, hyp: IterableSimpleNamespace, stretch: bool = False):
-    """
-    Apply a series of image transformations for training.
+    """Apply 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.
+    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 (IterableSimpleNamespace): A dictionary of hyperparameters controlling various aspects of the transformations.
+        hyp (IterableSimpleNamespace): 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:
@@ -2550,6 +2392,12 @@ def v8_transforms(dataset, imgsz: int, hyp: IterableSimpleNamespace, stretch: bo
         >>> 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])
+
+        >>> # With custom albumentations
+        >>> import albumentations as A
+        >>> augmentations = [A.Blur(p=0.01), A.CLAHE(p=0.01)]
+        >>> hyp.augmentations = augmentations
+        >>> transforms = v8_transforms(dataset, imgsz=640, hyp=hyp)
     """
     mosaic = Mosaic(dataset, imgsz=imgsz, p=hyp.mosaic)
     affine = RandomPerspective(
@@ -2587,7 +2435,7 @@ def v8_transforms(dataset, imgsz: int, hyp: IterableSimpleNamespace, stretch: bo
             pre_transform,
             MixUp(dataset, pre_transform=pre_transform, p=hyp.mixup),
             CutMix(dataset, pre_transform=pre_transform, p=hyp.cutmix),
-            Albumentations(p=1.0),
+            Albumentations(p=1.0, transforms=getattr(hyp, "augmentations", None)),
             RandomHSV(hgain=hyp.hsv_h, sgain=hyp.hsv_s, vgain=hyp.hsv_v),
             RandomFlip(direction="vertical", p=hyp.flipud, flip_idx=flip_idx),
             RandomFlip(direction="horizontal", p=hyp.fliplr, flip_idx=flip_idx),
@@ -2601,14 +2449,13 @@ def classify_transforms(
     mean: tuple[float, float, float] = DEFAULT_MEAN,
     std: tuple[float, float, float] = DEFAULT_STD,
     interpolation: str = "BILINEAR",
-    crop_fraction: float = None,
+    crop_fraction: float | None = None,
 ):
-    """
-    Create a composition of image transforms for classification tasks.
+    """Create 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.
+    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 | tuple): The target size for the transformed image. If an int, it defines the shortest edge. If a
@@ -2651,11 +2498,11 @@ def classify_augmentations(
     size: int = 224,
     mean: tuple[float, float, float] = DEFAULT_MEAN,
     std: tuple[float, float, float] = DEFAULT_STD,
-    scale: tuple[float, float] = None,
-    ratio: tuple[float, float] = None,
+    scale: tuple[float, float] | None = None,
+    ratio: tuple[float, float] | None = None,
     hflip: float = 0.5,
     vflip: float = 0.0,
-    auto_augment: str = None,
+    auto_augment: str | None = None,
     hsv_h: float = 0.015,  # image HSV-Hue augmentation (fraction)
     hsv_s: float = 0.4,  # image HSV-Saturation augmentation (fraction)
     hsv_v: float = 0.4,  # image HSV-Value augmentation (fraction)
@@ -2663,8 +2510,7 @@ def classify_augmentations(
     erasing: float = 0.0,
     interpolation: str = "BILINEAR",
 ):
-    """
-    Create a composition of image augmentation transforms for classification tasks.
+    """Create 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.
@@ -2752,11 +2598,10 @@ def classify_augmentations(
 
 # NOTE: keep this class for backward compatibility
 class ClassifyLetterBox:
-    """
-    A class for resizing and padding images for classification tasks.
+    """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.
+    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.
@@ -2776,30 +2621,16 @@ class ClassifyLetterBox:
     """
 
     def __init__(self, size: int | tuple[int, int] = (640, 640), auto: bool = False, stride: int = 32):
-        """
-        Initialize the ClassifyLetterBox object for image preprocessing.
+        """Initialize 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 (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).
+            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.
             stride (int): The stride value, used when 'auto' is True.
-
-        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
@@ -2807,8 +2638,7 @@ class ClassifyLetterBox:
         self.stride = stride  # used with auto
 
     def __call__(self, im: np.ndarray) -> np.ndarray:
-        """
-        Resize and pad an image using the letterbox method.
+        """Resize and pad 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.
@@ -2817,8 +2647,8 @@ class ClassifyLetterBox:
             im (np.ndarray): Input image as a numpy array with shape (H, W, C).
 
         Returns:
-            (np.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.
+            (np.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))
@@ -2843,8 +2673,7 @@ class ClassifyLetterBox:
 
 # NOTE: keep this class for backward compatibility
 class CenterCrop:
-    """
-    Apply center cropping to images for classification tasks.
+    """Apply 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()]).
@@ -2865,39 +2694,30 @@ class CenterCrop:
     """
 
     def __init__(self, size: int | tuple[int, int] = (640, 640)):
-        """
-        Initialize the CenterCrop object for image preprocessing.
+        """Initialize 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.
+            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: Image.Image | np.ndarray) -> np.ndarray:
-        """
-        Apply center cropping to an input image.
+        """Apply 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.
+        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 (np.ndarray | PIL.Image.Image): The input image as a numpy array of shape (H, W, C) or a
-                PIL Image object.
+            im (np.ndarray | PIL.Image.Image): The input image as a numpy array of shape (H, W, C) or a PIL Image
+                object.
 
         Returns:
             (np.ndarray): The center-cropped and resized image as a numpy array of shape (self.h, self.w, C).
@@ -2918,8 +2738,7 @@ class CenterCrop:
 
 # NOTE: keep this class for backward compatibility
 class ToTensor:
-    """
-    Convert an image from a numpy array to a PyTorch tensor.
+    """Convert 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()]).
 
@@ -2942,40 +2761,31 @@ class ToTensor:
     """
 
     def __init__(self, half: bool = False):
-        """
-        Initialize the ToTensor object for converting images to PyTorch tensors.
+        """Initialize 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.
+        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).
-
-        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: np.ndarray) -> torch.Tensor:
-        """
-        Transform an image from a numpy array to a PyTorch tensor.
+        """Transform 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.
+        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 (np.ndarray): Input image as a numpy array with shape (H, W, C) in RGB order.
 
         Returns:
-            (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.
+            (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)