Reference for ultralytics/data/augment.py

class BaseTransform:
    """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.

    Methods:
        apply_image: Apply image transformations to labels.
        apply_instances: Apply transformations to object instances in labels.
        apply_semantic: Apply semantic segmentation to an image.
        __call__: Apply 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:
        """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.
        """
        pass

def __call__(self, labels):
    """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.

    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)

def apply_image(self, 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.

    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):
    """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
    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):
    """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.

    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

class Compose:
    """A class for composing multiple image transformations.

    Attributes:
        transforms (list[Callable]): A list of transformation functions to be applied sequentially.

    Methods:
        __call__: Apply a series of transformations to input data.
        append: Append a new transform to the existing list of transforms.
        insert: Insert a new transform at a specified index in the list of transforms.
        __getitem__: Retrieve a specific transform or a set of transforms using indexing.
        __setitem__: Set a specific transform or a set of transforms using indexing.
        tolist: Convert 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):
        """Initialize the Compose object with a list of transforms.

        Args:
            transforms (list[Callable]): A list of callable transform objects to be applied sequentially.
        """
        self.transforms = transforms if isinstance(transforms, list) else [transforms]

def __call__(self, 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.

    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 __getitem__(self, index: list | int) -> Compose:
    """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.

    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)}"
    return Compose([self.transforms[i] for i in index]) if isinstance(index, list) else self.transforms[index]

def __repr__(self):
    """Return 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])})"

def __setitem__(self, index: list | int, value: list | int) -> None:
    """Set 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 append(self, transform):
    """Append 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):
    """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.
        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 tolist(self):
    """Convert 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

class BaseMixTransform:
    """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.

    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__: Apply 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: Update 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:
        """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.

        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].
        """
        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.

    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.

    Returns:
        (dict[str, Any]): 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

    # Get index of one or three other images
    indexes = self.get_indexes()
    if isinstance(indexes, int):
        indexes = [indexes]

    # Get images information will be used for Mosaic, CutMix or MixUp
    mix_labels = [self.dataset.get_image_and_label(i) for i in indexes]

    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, CutMix or MixUp
    labels = self._mix_transform(labels)
    labels.pop("mix_labels", None)
    return labels

def _mix_transform(self, labels: dict[str, Any]):
    """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.

    Args:
        labels (dict[str, Any]): 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[str, Any]): 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

@staticmethod
def _update_label_text(labels: dict[str, Any]) -> dict[str, Any]:
    """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.

    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.

    Returns:
        (dict[str, Any]): 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 = [*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)}

    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

def get_indexes(self):
    """Get 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]
    """
    return random.randint(0, len(self.dataset) - 1)

class Mosaic(BaseMixTransform):
    """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): 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: Return a list of random indexes from the dataset.
        _mix_transform: Apply mixup transformation to the input image and labels.
        _mosaic3: Create a 1x3 image mosaic.
        _mosaic4: Create a 2x2 image mosaic.
        _mosaic9: Create a 3x3 image mosaic.
        _update_labels: Update labels with padding.
        _cat_labels: Concatenate 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: int = 640, p: float = 1.0, n: int = 4):
        """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.

        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).
        """
        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."
        super().__init__(dataset=dataset, p=p)
        self.imgsz = imgsz
        self.border = (-imgsz // 2, -imgsz // 2)  # width, height
        self.n = n
        self.buffer_enabled = self.dataset.cache != "ram"

def _cat_labels(self, mosaic_labels: list[dict[str, Any]]) -> dict[str, Any]:
    """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.

    Args:
        mosaic_labels (list[dict[str, Any]]): A list of label dictionaries for each image in the mosaic.

    Returns:
        (dict[str, Any]): 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 not mosaic_labels:
        return {}
    cls = []
    instances = []
    imgsz = self.imgsz * 2  # mosaic imgsz
    for labels in mosaic_labels:
        cls.append(labels["cls"])
        instances.append(labels["instances"])
    # Final labels
    final_labels = {
        "im_file": mosaic_labels[0]["im_file"],
        "ori_shape": mosaic_labels[0]["ori_shape"],
        "resized_shape": (imgsz, imgsz),
        "cls": np.concatenate(cls, 0),
        "instances": Instances.concatenate(instances, axis=0),
        "mosaic_border": self.border,
    }
    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

def _mix_transform(self, labels: dict[str, Any]) -> dict[str, Any]:
    """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.

    Args:
        labels (dict[str, Any]): 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[str, Any]): 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") is None, "rect and mosaic are mutually exclusive."
    assert len(labels.get("mix_labels", [])), "There are no other images for mosaic augment."
    return (
        self._mosaic3(labels) if self.n == 3 else self._mosaic4(labels) if self.n == 4 else self._mosaic9(labels)
    )  # 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.

    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.

    Returns:
        (dict[str, Any]): 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):
        labels_patch = labels if i == 0 else labels["mix_labels"][i - 1]
        # Load image
        img = labels_patch["img"]
        h, w = labels_patch.pop("resized_shape")

        # Place img in img3
        if i == 0:  # center
            img3 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8)  # base image with 3 tiles
            h0, w0 = h, w
            c = s, s, s + w, s + h  # xmin, ymin, xmax, ymax (base) coordinates
        elif i == 1:  # right
            c = s + w0, s, s + w0 + w, s + h
        elif i == 2:  # left
            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 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

        # Labels assuming imgsz*2 mosaic size
        labels_patch = self._update_labels(labels_patch, padw + self.border[0], padh + self.border[1])
        mosaic_labels.append(labels_patch)
    final_labels = self._cat_labels(mosaic_labels)

    final_labels["img"] = img3[-self.border[0] : self.border[0], -self.border[1] : self.border[1]]
    return final_labels

def _mosaic4(self, labels: dict[str, Any]) -> dict[str, Any]:
    """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.

    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.

    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.

    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
    for i in range(4):
        labels_patch = labels if i == 0 else labels["mix_labels"][i - 1]
        # Load image
        img = labels_patch["img"]
        h, w = labels_patch.pop("resized_shape")

        # Place img in img4
        if i == 0:  # top left
            img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
            x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
            x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)
        elif i == 1:  # top right
            x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
            x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
        elif i == 2:  # bottom left
            x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
            x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
        elif i == 3:  # bottom right
            x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
            x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)

        img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
        padw = x1a - x1b
        padh = y1a - y1b

        labels_patch = self._update_labels(labels_patch, padw, padh)
        mosaic_labels.append(labels_patch)
    final_labels = self._cat_labels(mosaic_labels)
    final_labels["img"] = img4
    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.

    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:
            - '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.

    Returns:
        (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.

    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
    for i in range(9):
        labels_patch = labels if i == 0 else labels["mix_labels"][i - 1]
        # Load image
        img = labels_patch["img"]
        h, w = labels_patch.pop("resized_shape")

        # Place img in img9
        if i == 0:  # center
            img9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
            h0, w0 = h, w
            c = s, s, s + w, s + h  # xmin, ymin, xmax, ymax (base) coordinates
        elif i == 1:  # top
            c = s, s - h, s + w, s
        elif i == 2:  # top right
            c = s + wp, s - h, s + wp + w, s
        elif i == 3:  # right
            c = s + w0, s, s + w0 + w, s + h
        elif i == 4:  # bottom right
            c = s + w0, s + hp, s + w0 + w, s + hp + h
        elif i == 5:  # bottom
            c = s + w0 - w, s + h0, s + w0, s + h0 + h
        elif i == 6:  # bottom left
            c = s + w0 - wp - w, s + h0, s + w0 - wp, s + h0 + h
        elif i == 7:  # left
            c = s - w, s + h0 - h, s, s + h0
        elif i == 8:  # top left
            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 coordinates

        # Image
        img9[y1:y2, x1:x2] = img[y1 - padh :, x1 - padw :]  # img9[ymin:ymax, xmin:xmax]
        hp, wp = h, w  # height, width previous for next iteration

        # Labels assuming imgsz*2 mosaic size
        labels_patch = self._update_labels(labels_patch, padw + self.border[0], padh + self.border[1])
        mosaic_labels.append(labels_patch)
    final_labels = self._cat_labels(mosaic_labels)

    final_labels["img"] = img9[-self.border[0] : self.border[0], -self.border[1] : self.border[1]]
    return final_labels

@staticmethod
def _update_labels(labels, padw: int, padh: int) -> dict[str, Any]:
    """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.

    Args:
        labels (dict[str, Any]): 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)
    labels["instances"].add_padding(padw, padh)
    return labels

def get_indexes(self):
    """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.

    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 self.buffer_enabled:  # 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)]

class MixUp(BaseMixTransform):
    """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.

    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:
        _mix_transform: Apply 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: float = 0.0) -> None:
        """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.

        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].
        """
        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.

    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.

    Returns:
        (dict[str, Any]): 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)
    labels["instances"] = Instances.concatenate([labels["instances"], labels2["instances"]], axis=0)
    labels["cls"] = np.concatenate([labels["cls"], labels2["cls"]], 0)
    return labels

class CutMix(BaseMixTransform):
    """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.

    Attributes:
        dataset (Any): The dataset to which CutMix augmentation will be applied.
        pre_transform (Callable | None): Optional transform to apply before CutMix.
        p (float): Probability of applying CutMix augmentation.
        beta (float): Beta distribution parameter for sampling the mixing ratio.
        num_areas (int): Number of areas to try to cut and mix.

    Methods:
        _mix_transform: Apply CutMix augmentation to the input labels.
        _rand_bbox: Generate random bounding box coordinates for the cut region.

    Examples:
        >>> from ultralytics.data.augment import CutMix
        >>> dataset = YourDataset(...)  # Your image dataset
        >>> cutmix = CutMix(dataset, p=0.5)
        >>> augmented_labels = cutmix(original_labels)
    """

    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.

        Args:
            dataset (Any): The dataset to which CutMix augmentation will be applied.
            pre_transform (Callable | None): Optional transform to apply before CutMix.
            p (float): Probability of applying CutMix augmentation.
            beta (float): Beta distribution parameter for sampling the mixing ratio.
            num_areas (int): Number of areas to try to cut and mix.
        """
        super().__init__(dataset=dataset, pre_transform=pre_transform, p=p)
        self.beta = beta
        self.num_areas = num_areas