Reference for ultralytics/data/augment.py

def __call__(self, labels):
    """
    Applies all label transformations to an image, instances, and semantic masks.

    This method orchestrates the application of various transformations defined in the BaseTransform class
    to the input labels. It sequentially calls the apply_image and apply_instances methods to process the
    image and object instances, respectively.

    Args:
        labels (dict): A dictionary containing image data and annotations. Expected keys include 'img' for
            the image data, and 'instances' for object instances.

    Returns:
        (dict): The input labels dictionary with transformed image and instances.

    Examples:
        >>> transform = BaseTransform()
        >>> labels = {"img": np.random.rand(640, 640, 3), "instances": []}
        >>> transformed_labels = transform(labels)
    """
    self.apply_image(labels)
    self.apply_instances(labels)
    self.apply_semantic(labels)

def apply_image(self, labels):
    """
    Applies image transformations to labels.

    This method is intended to be overridden by subclasses to implement specific image transformation
    logic. In its base form, it returns the input labels unchanged.

    Args:
        labels (Any): The input labels to be transformed. The exact type and structure of labels may
            vary depending on the specific implementation.

    Returns:
        (Any): The transformed labels. In the base implementation, this is identical to the input.

    Examples:
        >>> transform = BaseTransform()
        >>> original_labels = [1, 2, 3]
        >>> transformed_labels = transform.apply_image(original_labels)
        >>> print(transformed_labels)
        [1, 2, 3]
    """
    pass

def apply_instances(self, labels):
    """
    Applies transformations to object instances in labels.

    This method is responsible for applying various transformations to object instances within the given
    labels. It is designed to be overridden by subclasses to implement specific instance transformation
    logic.

    Args:
        labels (dict): A dictionary containing label information, including object instances.

    Returns:
        (dict): The modified labels dictionary with transformed object instances.

    Examples:
        >>> transform = BaseTransform()
        >>> labels = {"instances": Instances(xyxy=torch.rand(5, 4), cls=torch.randint(0, 80, (5,)))}
        >>> transformed_labels = transform.apply_instances(labels)
    """
    pass

def apply_semantic(self, labels):
    """
    Applies semantic segmentation transformations to an image.

    This method is intended to be overridden by subclasses to implement specific semantic segmentation
    transformations. In its base form, it does not perform any operations.

    Args:
        labels (Any): The input labels or semantic segmentation mask to be transformed.

    Returns:
        (Any): The transformed semantic segmentation mask or labels.

    Examples:
        >>> transform = BaseTransform()
        >>> semantic_mask = np.zeros((100, 100), dtype=np.uint8)
        >>> transformed_mask = transform.apply_semantic(semantic_mask)
    """
    pass

def __call__(self, data):
    """
    Applies a series of transformations to input data. This method sequentially applies each transformation in the
    Compose object's list of transforms to the input data.

    Args:
        data (Any): The input data to be transformed. This can be of any type, depending on the
            transformations in the list.

    Returns:
        (Any): The transformed data after applying all transformations in sequence.

    Examples:
        >>> transforms = [Transform1(), Transform2(), Transform3()]
        >>> compose = Compose(transforms)
        >>> transformed_data = compose(input_data)
    """
    for t in self.transforms:
        data = t(data)
    return data

def __getitem__(self, index: Union[list, int]) -> "Compose":
    """
    Retrieves a specific transform or a set of transforms using indexing.

    Args:
        index (int | List[int]): Index or list of indices of the transforms to retrieve.

    Returns:
        (Compose): A new Compose object containing the selected transform(s).

    Raises:
        AssertionError: If the index is not of type int or list.

    Examples:
        >>> transforms = [RandomFlip(), RandomPerspective(10), RandomHSV(0.5, 0.5, 0.5)]
        >>> compose = Compose(transforms)
        >>> single_transform = compose[1]  # Returns a Compose object with only RandomPerspective
        >>> multiple_transforms = compose[0:2]  # Returns a Compose object with RandomFlip and RandomPerspective
    """
    assert isinstance(index, (int, list)), f"The indices should be either list or int type but got {type(index)}"
    index = [index] if isinstance(index, int) else index
    return Compose([self.transforms[i] for i in index])

def __repr__(self):
    """
    Returns a string representation of the Compose object.

    Returns:
        (str): A string representation of the Compose object, including the list of transforms.

    Examples:
        >>> transforms = [RandomFlip(), RandomPerspective(degrees=10, translate=0.1, scale=0.1)]
        >>> compose = Compose(transforms)
        >>> print(compose)
        Compose([
            RandomFlip(),
            RandomPerspective(degrees=10, translate=0.1, scale=0.1)
        ])
    """
    return f"{self.__class__.__name__}({', '.join([f'{t}' for t in self.transforms])})"

def __setitem__(self, index: Union[list, int], value: Union[list, int]) -> None:
    """
    Sets one or more transforms in the composition using indexing.

    Args:
        index (int | List[int]): Index or list of indices to set transforms at.
        value (Any | List[Any]): Transform or list of transforms to set at the specified index(es).

    Raises:
        AssertionError: If index type is invalid, value type doesn't match index type, or index is out of range.

    Examples:
        >>> compose = Compose([Transform1(), Transform2(), Transform3()])
        >>> compose[1] = NewTransform()  # Replace second transform
        >>> compose[0:2] = [NewTransform1(), NewTransform2()]  # Replace first two transforms
    """
    assert isinstance(index, (int, list)), f"The indices should be either list or int type but got {type(index)}"
    if isinstance(index, list):
        assert isinstance(value, list), (
            f"The indices should be the same type as values, but got {type(index)} and {type(value)}"
        )
    if isinstance(index, int):
        index, value = [index], [value]
    for i, v in zip(index, value):
        assert i < len(self.transforms), f"list index {i} out of range {len(self.transforms)}."
        self.transforms[i] = v

def append(self, transform):
    """
    Appends a new transform to the existing list of transforms.

    Args:
        transform (BaseTransform): The transformation to be added to the composition.

    Examples:
        >>> compose = Compose([RandomFlip(), RandomPerspective()])
        >>> compose.append(RandomHSV())
    """
    self.transforms.append(transform)

def insert(self, index, transform):
    """
    Inserts a new transform at a specified index in the existing list of transforms.

    Args:
        index (int): The index at which to insert the new transform.
        transform (BaseTransform): The transform object to be inserted.

    Examples:
        >>> compose = Compose([Transform1(), Transform2()])
        >>> compose.insert(1, Transform3())
        >>> len(compose.transforms)
        3
    """
    self.transforms.insert(index, transform)

def tolist(self):
    """
    Converts the list of transforms to a standard Python list.

    Returns:
        (list): A list containing all the transform objects in the Compose instance.

    Examples:
        >>> transforms = [RandomFlip(), RandomPerspective(10), CenterCrop()]
        >>> compose = Compose(transforms)
        >>> transform_list = compose.tolist()
        >>> print(len(transform_list))
        3
    """
    return self.transforms

def __call__(self, labels):
    """
    Applies pre-processing transforms and mixup/mosaic transforms to labels data.

    This method determines whether to apply the mix transform based on a probability factor. If applied, it
    selects additional images, applies pre-transforms if specified, and then performs the mix transform.

    Args:
        labels (dict): A dictionary containing label data for an image.

    Returns:
        (dict): The transformed labels dictionary, which may include mixed data from other images.

    Examples:
        >>> transform = BaseMixTransform(dataset, pre_transform=None, p=0.5)
        >>> result = transform({"image": img, "bboxes": boxes, "cls": classes})
    """
    if random.uniform(0, 1) > self.p:
        return labels

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

def get_indexes(self):
    """
    Gets a list of shuffled indexes for mosaic augmentation.

    Returns:
        (List[int]): A list of shuffled indexes from the dataset.

    Examples:
        >>> transform = BaseMixTransform(dataset)
        >>> indexes = transform.get_indexes()
        >>> print(indexes)  # [3, 18, 7, 2]
    """
    raise NotImplementedError

def get_indexes(self, buffer=True):
    """
    Returns a list of random indexes from the dataset for mosaic augmentation.

    This method selects random image indexes either from a buffer or from the entire dataset, depending on
    the 'buffer' parameter. It is used to choose images for creating mosaic augmentations.

    Args:
        buffer (bool): If True, selects images from the dataset buffer. If False, selects from the entire
            dataset.

    Returns:
        (List[int]): A list of random image indexes. The length of the list is n-1, where n is the number
            of images used in the mosaic (either 3 or 8, depending on whether n is 4 or 9).

    Examples:
        >>> mosaic = Mosaic(dataset, imgsz=640, p=1.0, n=4)
        >>> indexes = mosaic.get_indexes()
        >>> print(len(indexes))  # Output: 3
    """
    if buffer:  # select images from buffer
        return random.choices(list(self.dataset.buffer), k=self.n - 1)
    else:  # select any images
        return [random.randint(0, len(self.dataset) - 1) for _ in range(self.n - 1)]

def get_indexes(self):
    """
    Get a random index from the dataset.

    This method returns a single random index from the dataset, which is used to select an image for MixUp
    augmentation.

    Returns:
        (int): A random integer index within the range of the dataset length.

    Examples:
        >>> mixup = MixUp(dataset)
        >>> index = mixup.get_indexes()
        >>> print(index)
        42
    """
    return random.randint(0, len(self.dataset) - 1)

def __call__(self, labels):
    """
    Applies random perspective and affine transformations to an image and its associated labels.

    This method performs a series of transformations including rotation, translation, scaling, shearing,
    and perspective distortion on the input image and adjusts the corresponding bounding boxes, segments,
    and keypoints accordingly.

    Args:
        labels (dict): A 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): Transformed labels dictionary containing:
            - 'img' (np.ndarray): The transformed image.
            - 'cls' (np.ndarray): Updated class labels.
            - 'instances' (Instances): Updated object instances.
            - 'resized_shape' (Tuple[int, int]): New image shape after transformation.

    Examples:
        >>> transform = RandomPerspective()
        >>> image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
        >>> labels = {
        ...     "img": image,
        ...     "cls": np.array([0, 1, 2]),
        ...     "instances": Instances(bboxes=np.array([[10, 10, 50, 50], [100, 100, 150, 150]])),
        ... }
        >>> result = transform(labels)
        >>> assert result["img"].shape[:2] == result["resized_shape"]
    """
    if self.pre_transform and "mosaic_border" not in labels:
        labels = self.pre_transform(labels)
    labels.pop("ratio_pad", None)  # do not need ratio pad

    img = labels["img"]
    cls = labels["cls"]
    instances = labels.pop("instances")
    # Make sure the coord formats are right
    instances.convert_bbox(format="xyxy")
    instances.denormalize(*img.shape[:2][::-1])

    border = labels.pop("mosaic_border", self.border)
    self.size = img.shape[1] + border[1] * 2, img.shape[0] + border[0] * 2  # w, h
    # M is affine matrix
    # Scale for func:`box_candidates`
    img, M, scale = self.affine_transform(img, border)

    bboxes = self.apply_bboxes(instances.bboxes, M)

    segments = instances.segments
    keypoints = instances.keypoints
    # Update bboxes if there are segments.
    if len(segments):
        bboxes, segments = self.apply_segments(segments, M)

    if keypoints is not None:
        keypoints = self.apply_keypoints(keypoints, M)
    new_instances = Instances(bboxes, segments, keypoints, bbox_format="xyxy", normalized=False)
    # Clip
    new_instances.clip(*self.size)

    # Filter instances
    instances.scale(scale_w=scale, scale_h=scale, bbox_only=True)
    # Make the bboxes have the same scale with new_bboxes
    i = self.box_candidates(
        box1=instances.bboxes.T, box2=new_instances.bboxes.T, area_thr=0.01 if len(segments) else 0.10
    )
    labels["instances"] = new_instances[i]
    labels["cls"] = cls[i]
    labels["img"] = img
    labels["resized_shape"] = img.shape[:2]
    return labels

def affine_transform(self, img, border):
    """
    Applies a sequence of affine transformations centered around the image center.

    This function performs a series of geometric transformations on the input image, including
    translation, perspective change, rotation, scaling, and shearing. The transformations are
    applied in a specific order to maintain consistency.

    Args:
        img (np.ndarray): Input image to be transformed.
        border (Tuple[int, int]): Border dimensions for the transformed image.

    Returns:
        (Tuple[np.ndarray, np.ndarray, float]): A tuple containing:
            - np.ndarray: Transformed image.
            - np.ndarray: 3x3 transformation matrix.
            - float: Scale factor applied during the transformation.

    Examples:
        >>> import numpy as np
        >>> img = np.random.rand(100, 100, 3)
        >>> border = (10, 10)
        >>> transformed_img, matrix, scale = affine_transform(img, border)
    """
    # Center
    C = np.eye(3, dtype=np.float32)

    C[0, 2] = -img.shape[1] / 2  # x translation (pixels)
    C[1, 2] = -img.shape[0] / 2  # y translation (pixels)

    # Perspective
    P = np.eye(3, dtype=np.float32)
    P[2, 0] = random.uniform(-self.perspective, self.perspective)  # x perspective (about y)
    P[2, 1] = random.uniform(-self.perspective, self.perspective)  # y perspective (about x)

    # Rotation and Scale
    R = np.eye(3, dtype=np.float32)
    a = random.uniform(-self.degrees, self.degrees)
    # a += random.choice([-180, -90, 0, 90])  # add 90deg rotations to small rotations
    s = random.uniform(1 - self.scale, 1 + self.scale)
    # s = 2 ** random.uniform(-scale, scale)
    R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)

    # Shear
    S = np.eye(3, dtype=np.float32)
    S[0, 1] = math.tan(random.uniform(-self.shear, self.shear) * math.pi / 180)  # x shear (deg)
    S[1, 0] = math.tan(random.uniform(-self.shear, self.shear) * math.pi / 180)  # y shear (deg)

    # Translation
    T = np.eye(3, dtype=np.float32)
    T[0, 2] = random.uniform(0.5 - self.translate, 0.5 + self.translate) * self.size[0]  # x translation (pixels)
    T[1, 2] = random.uniform(0.5 - self.translate, 0.5 + self.translate) * self.size[1]  # y translation (pixels)

    # Combined rotation matrix
    M = T @ S @ R @ P @ C  # order of operations (right to left) is IMPORTANT
    # Affine image
    if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any():  # image changed
        if self.perspective:
            img = cv2.warpPerspective(img, M, dsize=self.size, borderValue=(114, 114, 114))
        else:  # affine
            img = cv2.warpAffine(img, M[:2], dsize=self.size, borderValue=(114, 114, 114))
    return img, M, s

def apply_bboxes(self, bboxes, M):
    """
    Apply affine transformation to bounding boxes.

    This function applies an affine transformation to a set of bounding boxes using the provided
    transformation matrix.

    Args:
        bboxes (torch.Tensor): Bounding boxes in xyxy format with shape (N, 4), where N is the number
            of bounding boxes.
        M (torch.Tensor): Affine transformation matrix with shape (3, 3).

    Returns:
        (torch.Tensor): Transformed bounding boxes in xyxy format with shape (N, 4).

    Examples:
        >>> bboxes = torch.tensor([[10, 10, 20, 20], [30, 30, 40, 40]])
        >>> M = torch.eye(3)
        >>> transformed_bboxes = apply_bboxes(bboxes, M)
    """
    n = len(bboxes)
    if n == 0:
        return bboxes

    xy = np.ones((n * 4, 3), dtype=bboxes.dtype)
    xy[:, :2] = bboxes[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(n * 4, 2)  # x1y1, x2y2, x1y2, x2y1
    xy = xy @ M.T  # transform
    xy = (xy[:, :2] / xy[:, 2:3] if self.perspective else xy[:, :2]).reshape(n, 8)  # perspective rescale or affine

    # Create new boxes
    x = xy[:, [0, 2, 4, 6]]
    y = xy[:, [1, 3, 5, 7]]
    return np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1)), dtype=bboxes.dtype).reshape(4, n).T

def apply_keypoints(self, keypoints, M):
    """
    Applies affine transformation to keypoints.

    This method transforms the input keypoints using the provided affine transformation matrix. It handles
    perspective rescaling if necessary and updates the visibility of keypoints that fall outside the image
    boundaries after transformation.

    Args:
        keypoints (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:
        (np.ndarray): Transformed keypoints array with the same shape as input (N, 17, 3).

    Examples:
        >>> random_perspective = RandomPerspective()
        >>> keypoints = np.random.rand(5, 17, 3)  # 5 instances, 17 keypoints each
        >>> M = np.eye(3)  # Identity transformation
        >>> transformed_keypoints = random_perspective.apply_keypoints(keypoints, M)
    """
    n, nkpt = keypoints.shape[:2]
    if n == 0:
        return keypoints
    xy = np.ones((n * nkpt, 3), dtype=keypoints.dtype)
    visible = keypoints[..., 2].reshape(n * nkpt, 1)
    xy[:, :2] = keypoints[..., :2].reshape(n * nkpt, 2)
    xy = xy @ M.T  # transform
    xy = xy[:, :2] / xy[:, 2:3]  # perspective rescale or affine
    out_mask = (xy[:, 0] < 0) | (xy[:, 1] < 0) | (xy[:, 0] > self.size[0]) | (xy[:, 1] > self.size[1])
    visible[out_mask] = 0
    return np.concatenate([xy, visible], axis=-1).reshape(n, nkpt, 3)

def apply_segments(self, segments, M):
    """
    Apply affine transformations to segments and generate new bounding boxes.

    This function applies affine transformations to input segments and generates new bounding boxes based on
    the transformed segments. It clips the transformed segments to fit within the new bounding boxes.

    Args:
        segments (np.ndarray): Input segments with shape (N, M, 2), where N is the number of segments and M is the
            number of points in each segment.
        M (np.ndarray): Affine transformation matrix with shape (3, 3).

    Returns:
        (Tuple[np.ndarray, np.ndarray]): A tuple containing:
            - New bounding boxes with shape (N, 4) in xyxy format.
            - Transformed and clipped segments with shape (N, M, 2).

    Examples:
        >>> segments = np.random.rand(10, 500, 2)  # 10 segments with 500 points each
        >>> M = np.eye(3)  # Identity transformation matrix
        >>> new_bboxes, new_segments = apply_segments(segments, M)
    """
    n, num = segments.shape[:2]
    if n == 0:
        return [], segments

    xy = np.ones((n * num, 3), dtype=segments.dtype)
    segments = segments.reshape(-1, 2)
    xy[:, :2] = segments
    xy = xy @ M.T  # transform
    xy = xy[:, :2] / xy[:, 2:3]
    segments = xy.reshape(n, -1, 2)
    bboxes = np.stack([segment2box(xy, self.size[0], self.size[1]) for xy in segments], 0)
    segments[..., 0] = segments[..., 0].clip(bboxes[:, 0:1], bboxes[:, 2:3])
    segments[..., 1] = segments[..., 1].clip(bboxes[:, 1:2], bboxes[:, 3:4])
    return bboxes, segments