Reference for ultralytics/data/augment.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/data/augment.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!
ultralytics.data.augment.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:
Name | Description |
---|---|
apply_image |
Applies image transformations to labels. |
apply_instances |
Applies transformations to object instances in labels. |
apply_semantic |
Applies semantic segmentation to an image. |
__call__ |
Applies all label transformations to an image, instances, and semantic masks. |
Examples:
>>> transform = BaseTransform()
>>> labels = {"image": np.array(...), "instances": [...], "semantic": np.array(...)}
>>> transformed_labels = transform(labels)
This constructor sets up the base transformation object, which can be extended for specific image processing tasks. It is designed to be compatible with both classification and semantic segmentation.
Examples:
Source code in ultralytics/data/augment.py
__call__
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict
|
A dictionary containing image data and annotations. Expected keys include 'img' for the image data, and 'instances' for object instances. |
required |
Returns:
Type | Description |
---|---|
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)
Source code in ultralytics/data/augment.py
apply_image
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Any
|
The input labels to be transformed. The exact type and structure of labels may vary depending on the specific implementation. |
required |
Returns:
Type | Description |
---|---|
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]
Source code in ultralytics/data/augment.py
apply_instances
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict
|
A dictionary containing label information, including object instances. |
required |
Returns:
Type | Description |
---|---|
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)
Source code in ultralytics/data/augment.py
apply_semantic
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Any
|
The input labels or semantic segmentation mask to be transformed. |
required |
Returns:
Type | Description |
---|---|
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)
Source code in ultralytics/data/augment.py
ultralytics.data.augment.Compose
A class for composing multiple image transformations.
Attributes:
Name | Type | Description |
---|---|---|
transforms |
List[Callable]
|
A list of transformation functions to be applied sequentially. |
Methods:
Name | Description |
---|---|
__call__ |
Applies a series of transformations to input data. |
append |
Appends a new transform to the existing list of transforms. |
insert |
Inserts a new transform at a specified index in the list of transforms. |
__getitem__ |
Retrieves a specific transform or a set of transforms using indexing. |
__setitem__ |
Sets a specific transform or a set of transforms using indexing. |
tolist |
Converts the list of transforms to a standard Python list. |
Examples:
>>> transforms = [RandomFlip(), RandomPerspective(30)]
>>> compose = Compose(transforms)
>>> transformed_data = compose(data)
>>> compose.append(CenterCrop((224, 224)))
>>> compose.insert(0, RandomFlip())
Parameters:
Name | Type | Description | Default |
---|---|---|---|
transforms
|
List[Callable]
|
A list of callable transform objects to be applied sequentially. |
required |
Examples:
>>> from ultralytics.data.augment import Compose, RandomHSV, RandomFlip
>>> transforms = [RandomHSV(), RandomFlip()]
>>> compose = Compose(transforms)
Source code in ultralytics/data/augment.py
__call__
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
data
|
Any
|
The input data to be transformed. This can be of any type, depending on the transformations in the list. |
required |
Returns:
Type | Description |
---|---|
Any
|
The transformed data after applying all transformations in sequence. |
Examples:
>>> transforms = [Transform1(), Transform2(), Transform3()]
>>> compose = Compose(transforms)
>>> transformed_data = compose(input_data)
Source code in ultralytics/data/augment.py
__getitem__
Retrieves a specific transform or a set of transforms using indexing.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
index
|
int | List[int]
|
Index or list of indices of the transforms to retrieve. |
required |
Returns:
Type | Description |
---|---|
Compose
|
A new Compose object containing the selected transform(s). |
Raises:
Type | Description |
---|---|
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
Source code in ultralytics/data/augment.py
__repr__
Returns a string representation of the Compose object.
Returns:
Type | Description |
---|---|
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)
])
Source code in ultralytics/data/augment.py
__setitem__
Sets one or more transforms in the composition using indexing.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
index
|
int | List[int]
|
Index or list of indices to set transforms at. |
required |
value
|
Any | List[Any]
|
Transform or list of transforms to set at the specified index(es). |
required |
Raises:
Type | Description |
---|---|
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
Source code in ultralytics/data/augment.py
append
Appends a new transform to the existing list of transforms.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
transform
|
BaseTransform
|
The transformation to be added to the composition. |
required |
Examples:
Source code in ultralytics/data/augment.py
insert
Inserts a new transform at a specified index in the existing list of transforms.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
index
|
int
|
The index at which to insert the new transform. |
required |
transform
|
BaseTransform
|
The transform object to be inserted. |
required |
Examples:
>>> compose = Compose([Transform1(), Transform2()])
>>> compose.insert(1, Transform3())
>>> len(compose.transforms)
3
Source code in ultralytics/data/augment.py
tolist
Converts the list of transforms to a standard Python list.
Returns:
Type | Description |
---|---|
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
Source code in ultralytics/data/augment.py
ultralytics.data.augment.BaseMixTransform
Base class for mix transformations like MixUp and Mosaic.
This class provides a foundation for implementing mix transformations on datasets. It handles the probability-based application of transforms and manages the mixing of multiple images and labels.
Attributes:
Name | Type | Description |
---|---|---|
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:
Name | Description |
---|---|
__call__ |
Applies the mix transformation to the input labels. |
_mix_transform |
Abstract method to be implemented by subclasses for specific mix operations. |
get_indexes |
Abstract method to get indexes of images to be mixed. |
_update_label_text |
Updates label text for mixed images. |
Examples:
>>> class CustomMixTransform(BaseMixTransform):
... def _mix_transform(self, labels):
... # Implement custom mix logic here
... return labels
...
... def get_indexes(self):
... return [random.randint(0, len(self.dataset) - 1) for _ in range(3)]
>>> dataset = YourDataset()
>>> transform = CustomMixTransform(dataset, p=0.5)
>>> mixed_labels = transform(original_labels)
This class serves as a base for implementing mix transformations in image processing pipelines.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset
|
Any
|
The dataset object containing images and labels for mixing. |
required |
pre_transform
|
Callable | None
|
Optional transform to apply before mixing. |
None
|
p
|
float
|
Probability of applying the mix transformation. Should be in the range [0.0, 1.0]. |
0.0
|
Examples:
>>> dataset = YOLODataset("path/to/data")
>>> pre_transform = Compose([RandomFlip(), RandomPerspective()])
>>> mix_transform = BaseMixTransform(dataset, pre_transform, p=0.5)
Source code in ultralytics/data/augment.py
__call__
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict
|
A dictionary containing label data for an image. |
required |
Returns:
Type | Description |
---|---|
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})
Source code in ultralytics/data/augment.py
get_indexes
Gets a list of shuffled indexes for mosaic augmentation.
Returns:
Type | Description |
---|---|
List[int]
|
A list of shuffled indexes from the dataset. |
Examples:
>>> transform = BaseMixTransform(dataset)
>>> indexes = transform.get_indexes()
>>> print(indexes) # [3, 18, 7, 2]
Source code in ultralytics/data/augment.py
ultralytics.data.augment.Mosaic
Bases: 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:
Name | Type | Description |
---|---|---|
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:
Name | Description |
---|---|
get_indexes |
Returns a list of random indexes from the dataset. |
_mix_transform |
Applies mixup transformation to the input image and labels. |
_mosaic3 |
Creates a 1x3 image mosaic. |
_mosaic4 |
Creates a 2x2 image mosaic. |
_mosaic9 |
Creates a 3x3 image mosaic. |
_update_labels |
Updates labels with padding. |
_cat_labels |
Concatenates labels and clips mosaic border instances. |
Examples:
>>> from ultralytics.data.augment import Mosaic
>>> dataset = YourDataset(...) # Your image dataset
>>> mosaic_aug = Mosaic(dataset, imgsz=640, p=0.5, n=4)
>>> augmented_labels = mosaic_aug(original_labels)
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset
|
Any
|
The dataset on which the mosaic augmentation is applied. |
required |
imgsz
|
int
|
Image size (height and width) after mosaic pipeline of a single image. |
640
|
p
|
float
|
Probability of applying the mosaic augmentation. Must be in the range 0-1. |
1.0
|
n
|
int
|
The grid size, either 4 (for 2x2) or 9 (for 3x3). |
4
|
Examples:
>>> from ultralytics.data.augment import Mosaic
>>> dataset = YourDataset(...)
>>> mosaic_aug = Mosaic(dataset, imgsz=640, p=0.5, n=4)
Source code in ultralytics/data/augment.py
get_indexes
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
buffer
|
bool
|
If True, selects images from the dataset buffer. If False, selects from the entire dataset. |
True
|
Returns:
Type | Description |
---|---|
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
Source code in ultralytics/data/augment.py
ultralytics.data.augment.MixUp
Bases: BaseMixTransform
Applies MixUp augmentation to image datasets.
This class implements the MixUp augmentation technique as described in the paper "mixup: Beyond Empirical Risk Minimization" (https://arxiv.org/abs/1710.09412). MixUp combines two images and their labels using a random weight.
Attributes:
Name | Type | Description |
---|---|---|
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:
Name | Description |
---|---|
get_indexes |
Returns a random index from the dataset. |
_mix_transform |
Applies MixUp augmentation to the input labels. |
Examples:
>>> from ultralytics.data.augment import MixUp
>>> dataset = YourDataset(...) # Your image dataset
>>> mixup = MixUp(dataset, p=0.5)
>>> augmented_labels = mixup(original_labels)
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset
|
Any
|
The dataset to which MixUp augmentation will be applied. |
required |
pre_transform
|
Callable | None
|
Optional transform to apply to images before MixUp. |
None
|
p
|
float
|
Probability of applying MixUp augmentation to an image. Must be in the range [0, 1]. |
0.0
|
Examples:
>>> from ultralytics.data.dataset import YOLODataset
>>> dataset = YOLODataset("path/to/data.yaml")
>>> mixup = MixUp(dataset, pre_transform=None, p=0.5)
Source code in ultralytics/data/augment.py
get_indexes
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:
Type | Description |
---|---|
int
|
A random integer index within the range of the dataset length. |
Examples:
Source code in ultralytics/data/augment.py
ultralytics.data.augment.RandomPerspective
RandomPerspective(
degrees=0.0,
translate=0.1,
scale=0.5,
shear=0.0,
perspective=0.0,
border=(0, 0),
pre_transform=None,
)
Implements random perspective and affine transformations on images and corresponding annotations.
This class applies random rotations, translations, scaling, shearing, and perspective transformations to images and their associated bounding boxes, segments, and keypoints. It can be used as part of an augmentation pipeline for object detection and instance segmentation tasks.
Attributes:
Name | Type | Description |
---|---|---|
degrees |
float
|
Maximum absolute degree range for random rotations. |
translate |
float
|
Maximum translation as a fraction of the image size. |
scale |
float
|
Scaling factor range, e.g., scale=0.1 means 0.9-1.1. |
shear |
float
|
Maximum shear angle in degrees. |
perspective |
float
|
Perspective distortion factor. |
border |
Tuple[int, int]
|
Mosaic border size as (x, y). |
pre_transform |
Callable | None
|
Optional transform to apply before the random perspective. |
Methods:
Name | Description |
---|---|
affine_transform |
Applies affine transformations to the input image. |
apply_bboxes |
Transforms bounding boxes using the affine matrix. |
apply_segments |
Transforms segments and generates new bounding boxes. |
apply_keypoints |
Transforms keypoints using the affine matrix. |
__call__ |
Applies the random perspective transformation to images and annotations. |
box_candidates |
Filters transformed bounding boxes based on size and aspect ratio. |
Examples:
>>> transform = RandomPerspective(degrees=10, translate=0.1, scale=0.1, shear=10)
>>> image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
>>> labels = {"img": image, "cls": np.array([0, 1]), "instances": Instances(...)}
>>> result = transform(labels)
>>> transformed_image = result["img"]
>>> transformed_instances = result["instances"]
This class implements random perspective and affine transformations on images and corresponding bounding boxes, segments, and keypoints. Transformations include rotation, translation, scaling, and shearing.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
degrees
|
float
|
Degree range for random rotations. |
0.0
|
translate
|
float
|
Fraction of total width and height for random translation. |
0.1
|
scale
|
float
|
Scaling factor interval, e.g., a scale factor of 0.5 allows a resize between 50%-150%. |
0.5
|
shear
|
float
|
Shear intensity (angle in degrees). |
0.0
|
perspective
|
float
|
Perspective distortion factor. |
0.0
|
border
|
Tuple[int, int]
|
Tuple specifying mosaic border (top/bottom, left/right). |
(0, 0)
|
pre_transform
|
Callable | None
|
Function/transform to apply to the image before starting the random transformation. |
None
|
Examples:
>>> transform = RandomPerspective(degrees=10.0, translate=0.1, scale=0.5, shear=5.0)
>>> result = transform(labels) # Apply random perspective to labels
Source code in ultralytics/data/augment.py
__call__
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict
|
A dictionary containing image data and annotations. Must include: 'img' (ndarray): The input image. 'cls' (ndarray): Class labels. 'instances' (Instances): Object instances with bounding boxes, segments, and keypoints. May include: 'mosaic_border' (Tuple[int, int]): Border size for mosaic augmentation. |
required |
Returns:
Type | Description |
---|---|
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"]
Source code in ultralytics/data/augment.py
1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 |
|
affine_transform
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
img
|
ndarray
|
Input image to be transformed. |
required |
border
|
Tuple[int, int]
|
Border dimensions for the transformed image. |
required |
Returns:
Type | Description |
---|---|
Tuple[ndarray, 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)
Source code in ultralytics/data/augment.py
1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 |
|
apply_bboxes
Apply affine transformation to bounding boxes.
This function applies an affine transformation to a set of bounding boxes using the provided transformation matrix.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
bboxes
|
Tensor
|
Bounding boxes in xyxy format with shape (N, 4), where N is the number of bounding boxes. |
required |
M
|
Tensor
|
Affine transformation matrix with shape (3, 3). |
required |
Returns:
Type | Description |
---|---|
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)
Source code in ultralytics/data/augment.py
apply_keypoints
Applies affine transformation to keypoints.
This method transforms the input keypoints using the provided affine transformation matrix. It handles perspective rescaling if necessary and updates the visibility of keypoints that fall outside the image boundaries after transformation.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
keypoints
|
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). |
required |
M
|
ndarray
|
3x3 affine transformation matrix. |
required |
Returns:
Type | Description |
---|---|
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)
Source code in ultralytics/data/augment.py
apply_segments
Apply affine transformations to segments and generate new bounding boxes.
This function applies affine transformations to input segments and generates new bounding boxes based on the transformed segments. It clips the transformed segments to fit within the new bounding boxes.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
segments
|
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. |
required |
M
|
ndarray
|
Affine transformation matrix with shape (3, 3). |
required |
Returns:
Type | Description |
---|---|
Tuple[ndarray, 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)
Source code in ultralytics/data/augment.py
box_candidates
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
box1
|
ndarray
|
Original boxes before augmentation, shape (4, N) where n is the number of boxes. Format is [x1, y1, x2, y2] in absolute coordinates. |
required |
box2
|
ndarray
|
Augmented boxes after transformation, shape (4, N). Format is [x1, y1, x2, y2] in absolute coordinates. |
required |
wh_thr
|
float
|
Width and height threshold in pixels. Boxes smaller than this in either dimension are rejected. |
2
|
ar_thr
|
float
|
Aspect ratio threshold. Boxes with an aspect ratio greater than this value are rejected. |
100
|
area_thr
|
float
|
Area ratio threshold. Boxes with an area ratio (new/old) less than this value are rejected. |
0.1
|
eps
|
float
|
Small epsilon value to prevent division by zero. |
1e-16
|
Returns:
Type | Description |
---|---|
ndarray
|
Boolean array of shape (n,) indicating which boxes are candidates. True values correspond to boxes that meet all criteria. |
Examples:
>>> random_perspective = RandomPerspective()
>>> box1 = np.array([[0, 0, 100, 100], [0, 0, 50, 50]]).T
>>> box2 = np.array([[10, 10, 90, 90], [5, 5, 45, 45]]).T
>>> candidates = random_perspective.box_candidates(box1, box2)
>>> print(candidates)
[True True]
Source code in ultralytics/data/augment.py
ultralytics.data.augment.RandomHSV
Randomly adjusts the Hue, Saturation, and Value (HSV) channels of an image.
This class applies random HSV augmentation to images within predefined limits set by hgain, sgain, and vgain.
Attributes:
Name | Type | Description |
---|---|---|
hgain |
float
|
Maximum variation for hue. Range is typically [0, 1]. |
sgain |
float
|
Maximum variation for saturation. Range is typically [0, 1]. |
vgain |
float
|
Maximum variation for value. Range is typically [0, 1]. |
Methods:
Name | Description |
---|---|
__call__ |
Applies random HSV augmentation to an image. |
Examples:
>>> import numpy as np
>>> from ultralytics.data.augment import RandomHSV
>>> augmenter = RandomHSV(hgain=0.5, sgain=0.5, vgain=0.5)
>>> image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
>>> labels = {"img": image}
>>> augmented_labels = augmenter(labels)
>>> augmented_image = augmented_labels["img"]
This class applies random adjustments to the HSV channels of an image within specified limits.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
hgain
|
float
|
Maximum variation for hue. Should be in the range [0, 1]. |
0.5
|
sgain
|
float
|
Maximum variation for saturation. Should be in the range [0, 1]. |
0.5
|
vgain
|
float
|
Maximum variation for value. Should be in the range [0, 1]. |
0.5
|
Examples:
Source code in ultralytics/data/augment.py
__call__
Applies 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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict
|
A dictionary containing image data and metadata. Must include an 'img' key with the image as a numpy array. |
required |
Returns:
Type | Description |
---|---|
None
|
The function modifies the input 'labels' dictionary in-place, updating the 'img' key with the HSV-augmented image. |
Examples:
>>> hsv_augmenter = RandomHSV(hgain=0.5, sgain=0.5, vgain=0.5)
>>> labels = {"img": np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)}
>>> hsv_augmenter(labels)
>>> augmented_img = labels["img"]
Source code in ultralytics/data/augment.py
ultralytics.data.augment.RandomFlip
Applies 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.
Attributes:
Name | Type | Description |
---|---|---|
p |
float
|
Probability of applying the flip. Must be between 0 and 1. |
direction |
str
|
Direction of flip, either 'horizontal' or 'vertical'. |
flip_idx |
array - like
|
Index mapping for flipping keypoints, if applicable. |
Methods:
Name | Description |
---|---|
__call__ |
Applies the random flip transformation to an image and its annotations. |
Examples:
>>> transform = RandomFlip(p=0.5, direction="horizontal")
>>> result = transform({"img": image, "instances": instances})
>>> flipped_image = result["img"]
>>> flipped_instances = result["instances"]
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
p
|
float
|
The probability of applying the flip. Must be between 0 and 1. |
0.5
|
direction
|
str
|
The direction to apply the flip. Must be 'horizontal' or 'vertical'. |
'horizontal'
|
flip_idx
|
List[int] | None
|
Index mapping for flipping keypoints, if any. |
None
|
Raises:
Type | Description |
---|---|
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 = RandomFlip(p=0.7, direction="vertical", flip_idx=[1, 0, 3, 2, 5, 4])
Source code in ultralytics/data/augment.py
__call__
Applies random flip to an image and updates any instances like bounding boxes or keypoints accordingly.
This method randomly flips the input image either horizontally or vertically based on the initialized probability and direction. It also updates the corresponding instances (bounding boxes, keypoints) to match the flipped image.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict
|
A dictionary containing the following keys: 'img' (numpy.ndarray): The image to be flipped. 'instances' (ultralytics.utils.instance.Instances): An object containing bounding boxes and optionally keypoints. |
required |
Returns:
Type | Description |
---|---|
Dict
|
The same dictionary with the flipped image and updated instances: 'img' (numpy.ndarray): The flipped image. 'instances' (ultralytics.utils.instance.Instances): Updated instances matching the flipped image. |
Examples:
>>> labels = {"img": np.random.rand(640, 640, 3), "instances": Instances(...)}
>>> random_flip = RandomFlip(p=0.5, direction="horizontal")
>>> flipped_labels = random_flip(labels)
Source code in ultralytics/data/augment.py
ultralytics.data.augment.LetterBox
LetterBox(
new_shape=(640, 640),
auto=False,
scaleFill=False,
scaleup=True,
center=True,
stride=32,
)
Resize image and padding for detection, instance segmentation, pose.
This class resizes and pads images to a specified shape while preserving aspect ratio. It also updates corresponding labels and bounding boxes.
Attributes:
Name | Type | Description |
---|---|---|
new_shape |
tuple
|
Target shape (height, width) for resizing. |
auto |
bool
|
Whether to use minimum rectangle. |
scaleFill |
bool
|
Whether to stretch the image to new_shape. |
scaleup |
bool
|
Whether to allow scaling up. If False, only scale down. |
stride |
int
|
Stride for rounding padding. |
center |
bool
|
Whether to center the image or align to top-left. |
Methods:
Name | Description |
---|---|
__call__ |
Resize and pad image, update labels and bounding boxes. |
Examples:
>>> transform = LetterBox(new_shape=(640, 640))
>>> result = transform(labels)
>>> resized_img = result["img"]
>>> updated_instances = result["instances"]
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
new_shape
|
Tuple[int, int]
|
Target size (height, width) for the resized image. |
(640, 640)
|
auto
|
bool
|
If True, use minimum rectangle to resize. If False, use new_shape directly. |
False
|
scaleFill
|
bool
|
If True, stretch the image to new_shape without padding. |
False
|
scaleup
|
bool
|
If True, allow scaling up. If False, only scale down. |
True
|
center
|
bool
|
If True, center the placed image. If False, place image in top-left corner. |
True
|
stride
|
int
|
Stride of the model (e.g., 32 for YOLOv5). |
32
|
Attributes:
Name | Type | Description |
---|---|---|
new_shape |
Tuple[int, int]
|
Target size for the resized image. |
auto |
bool
|
Flag for using minimum rectangle resizing. |
scaleFill |
bool
|
Flag for stretching image without padding. |
scaleup |
bool
|
Flag for allowing upscaling. |
stride |
int
|
Stride value for ensuring image size is divisible by stride. |
Examples:
>>> letterbox = LetterBox(new_shape=(640, 640), auto=False, scaleFill=False, scaleup=True, stride=32)
>>> resized_img = letterbox(original_img)
Source code in ultralytics/data/augment.py
__call__
Resizes and pads an image for object detection, instance segmentation, or pose estimation tasks.
This method applies letterboxing to the input image, which involves resizing the image while maintaining its aspect ratio and adding padding to fit the new shape. It also updates any associated labels accordingly.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict | None
|
A dictionary containing image data and associated labels, or empty dict if None. |
None
|
image
|
ndarray | None
|
The input image as a numpy array. If None, the image is taken from 'labels'. |
None
|
Returns:
Type | Description |
---|---|
Dict | Tuple
|
If 'labels' is provided, returns an updated dictionary with the resized and padded image, updated labels, and additional metadata. If 'labels' is empty, returns a tuple containing the resized and padded image, and a tuple of (ratio, (left_pad, top_pad)). |
Examples:
>>> letterbox = LetterBox(new_shape=(640, 640))
>>> result = letterbox(labels={"img": np.zeros((480, 640, 3)), "instances": Instances(...)})
>>> resized_img = result["img"]
>>> updated_instances = result["instances"]
Source code in ultralytics/data/augment.py
1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562 1563 1564 1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593 1594 1595 1596 1597 1598 1599 |
|
ultralytics.data.augment.CopyPaste
Bases: BaseMixTransform
CopyPaste class for applying Copy-Paste augmentation to image datasets.
This class implements the Copy-Paste augmentation technique as described in the paper "Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation" (https://arxiv.org/abs/2012.07177). It combines objects from different images to create new training samples.
Attributes:
Name | Type | Description |
---|---|---|
dataset |
Any
|
The dataset to which Copy-Paste augmentation will be applied. |
pre_transform |
Callable | None
|
Optional transform to apply before Copy-Paste. |
p |
float
|
Probability of applying Copy-Paste augmentation. |
Methods:
Name | Description |
---|---|
get_indexes |
Returns a random index from the dataset. |
_mix_transform |
Applies Copy-Paste augmentation to the input labels. |
__call__ |
Applies the Copy-Paste transformation to images and annotations. |
Examples:
>>> from ultralytics.data.augment import CopyPaste
>>> dataset = YourDataset(...) # Your image dataset
>>> copypaste = CopyPaste(dataset, p=0.5)
>>> augmented_labels = copypaste(original_labels)
Source code in ultralytics/data/augment.py
__call__
Applies Copy-Paste augmentation to an image and its labels.
Source code in ultralytics/data/augment.py
get_indexes
ultralytics.data.augment.Albumentations
Albumentations transformations for image augmentation.
This class applies various image transformations using the Albumentations library. It includes operations such as Blur, Median Blur, conversion to grayscale, Contrast Limited Adaptive Histogram Equalization (CLAHE), random changes in brightness and contrast, RandomGamma, and image quality reduction through compression.
Attributes:
Name | Type | Description |
---|---|---|
p |
float
|
Probability of applying the transformations. |
transform |
Compose
|
Composed Albumentations transforms. |
contains_spatial |
bool
|
Indicates if the transforms include spatial operations. |
Methods:
Name | Description |
---|---|
__call__ |
Applies the Albumentations transformations to the input labels. |
Examples:
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.
This class applies various image augmentations using the Albumentations library, including Blur, Median Blur, conversion to grayscale, Contrast Limited Adaptive Histogram Equalization, random changes of brightness and contrast, RandomGamma, and image quality reduction through compression.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
p
|
float
|
Probability of applying the augmentations. Must be between 0 and 1. |
1.0
|
Attributes:
Name | Type | Description |
---|---|---|
p |
float
|
Probability of applying the augmentations. |
transform |
Compose
|
Composed Albumentations transforms. |
contains_spatial |
bool
|
Indicates if the transforms include spatial transformations. |
Raises:
Type | Description |
---|---|
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.
Source code in ultralytics/data/augment.py
1758 1759 1760 1761 1762 1763 1764 1765 1766 1767 1768 1769 1770 1771 1772 1773 1774 1775 1776 1777 1778 1779 1780 1781 1782 1783 1784 1785 1786 1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811 1812 1813 1814 1815 1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832 1833 1834 1835 1836 1837 1838 1839 1840 1841 1842 1843 1844 1845 1846 1847 1848 1849 1850 1851 1852 1853 1854 1855 1856 1857 1858 1859 1860 1861 1862 1863 1864 |
|
__call__
Applies Albumentations transformations to input labels.
This method applies a series of image augmentations using the Albumentations library. It can perform both spatial and non-spatial transformations on the input image and its corresponding labels.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict
|
A dictionary containing image data and annotations. Expected keys are: - 'img': numpy.ndarray representing the image - 'cls': numpy.ndarray of class labels - 'instances': object containing bounding boxes and other instance information |
required |
Returns:
Type | Description |
---|---|
Dict
|
The input dictionary with augmented image and updated annotations. |
Examples:
>>> transform = Albumentations(p=0.5)
>>> labels = {
... "img": np.random.rand(640, 640, 3),
... "cls": np.array([0, 1]),
... "instances": Instances(bboxes=np.array([[0, 0, 1, 1], [0.5, 0.5, 0.8, 0.8]])),
... }
>>> augmented = transform(labels)
>>> assert augmented["img"].shape == (640, 640, 3)
Notes
- The method applies transformations with probability self.p.
- Spatial transforms update bounding boxes, while non-spatial transforms only modify the image.
- Requires the Albumentations library to be installed.
Source code in ultralytics/data/augment.py
ultralytics.data.augment.Format
Format(bbox_format='xywh', normalize=True, return_mask=False, return_keypoint=False, return_obb=False, mask_ratio=4, mask_overlap=True, batch_idx=True, bgr=0.0)
A class for formatting image annotations for object detection, instance segmentation, and pose estimation tasks.
This class standardizes image and instance annotations to be used by the collate_fn
in PyTorch DataLoader.
Attributes:
Name | Type | Description |
---|---|---|
bbox_format |
str
|
Format for bounding boxes. Options are 'xywh' or 'xyxy'. |
normalize |
bool
|
Whether to normalize bounding boxes. |
return_mask |
bool
|
Whether to return instance masks for segmentation. |
return_keypoint |
bool
|
Whether to return keypoints for pose estimation. |
return_obb |
bool
|
Whether to return oriented bounding boxes. |
mask_ratio |
int
|
Downsample ratio for masks. |
mask_overlap |
bool
|
Whether to overlap masks. |
batch_idx |
bool
|
Whether to keep batch indexes. |
bgr |
float
|
The probability to return BGR images. |
Methods:
Name | Description |
---|---|
__call__ |
Formats labels dictionary with image, classes, bounding boxes, and optionally masks and keypoints. |
_format_img |
Converts image from Numpy array to PyTorch tensor. |
_format_segments |
Converts polygon points to bitmap masks. |
Examples:
>>> formatter = Format(bbox_format="xywh", normalize=True, return_mask=True)
>>> formatted_labels = formatter(labels)
>>> img = formatted_labels["img"]
>>> bboxes = formatted_labels["bboxes"]
>>> masks = formatted_labels["masks"]
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
.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
bbox_format
|
str
|
Format for bounding boxes. Options are 'xywh', 'xyxy', etc. |
'xywh'
|
normalize
|
bool
|
Whether to normalize bounding boxes to [0,1]. |
True
|
return_mask
|
bool
|
If True, returns instance masks for segmentation tasks. |
False
|
return_keypoint
|
bool
|
If True, returns keypoints for pose estimation tasks. |
False
|
return_obb
|
bool
|
If True, returns oriented bounding boxes. |
False
|
mask_ratio
|
int
|
Downsample ratio for masks. |
4
|
mask_overlap
|
bool
|
If True, allows mask overlap. |
True
|
batch_idx
|
bool
|
If True, keeps batch indexes. |
True
|
bgr
|
float
|
Probability of returning BGR images instead of RGB. |
0.0
|
Attributes:
Name | Type | Description |
---|---|---|
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
Source code in ultralytics/data/augment.py
__call__
Formats image annotations for object detection, instance segmentation, and pose estimation tasks.
This method standardizes the image and instance annotations to be used by the collate_fn
in PyTorch
DataLoader. It processes the input labels dictionary, converting annotations to the specified format and
applying normalization if required.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict
|
A dictionary containing image and annotation data with the following keys: - 'img': The input image as a numpy array. - 'cls': Class labels for instances. - 'instances': An Instances object containing bounding boxes, segments, and keypoints. |
required |
Returns:
Type | Description |
---|---|
Dict
|
A dictionary with formatted data, including: - 'img': Formatted image tensor. - 'cls': Class labels tensor. - 'bboxes': Bounding boxes tensor in the specified format. - 'masks': Instance masks tensor (if return_mask is True). - 'keypoints': Keypoints tensor (if return_keypoint is True). - 'batch_idx': Batch index tensor (if batch_idx is True). |
Examples:
>>> formatter = Format(bbox_format="xywh", normalize=True, return_mask=True)
>>> labels = {"img": np.random.rand(640, 640, 3), "cls": np.array([0, 1]), "instances": Instances(...)}
>>> formatted_labels = formatter(labels)
>>> print(formatted_labels.keys())
Source code in ultralytics/data/augment.py
2005 2006 2007 2008 2009 2010 2011 2012 2013 2014 2015 2016 2017 2018 2019 2020 2021 2022 2023 2024 2025 2026 2027 2028 2029 2030 2031 2032 2033 2034 2035 2036 2037 2038 2039 2040 2041 2042 2043 2044 2045 2046 2047 2048 2049 2050 2051 2052 2053 2054 2055 2056 2057 2058 2059 2060 2061 2062 2063 2064 2065 2066 2067 2068 2069 2070 |
|
ultralytics.data.augment.RandomLoadText
RandomLoadText(prompt_format: str = '{}', neg_samples: Tuple[int, int] = (80, 80), max_samples: int = 80, padding: bool = False, padding_value: str = '')
Randomly samples positive and negative texts and updates class indices accordingly.
This class is responsible for sampling texts from a given set of class texts, including both positive (present in the image) and negative (not present in the image) samples. It updates the class indices to reflect the sampled texts and can optionally pad the text list to a fixed length.
Attributes:
Name | Type | Description |
---|---|---|
prompt_format |
str
|
Format string for text prompts. |
neg_samples |
Tuple[int, int]
|
Range for randomly sampling negative texts. |
max_samples |
int
|
Maximum number of different text samples in one image. |
padding |
bool
|
Whether to pad texts to max_samples. |
padding_value |
str
|
The text used for padding when padding is True. |
Methods:
Name | Description |
---|---|
__call__ |
Processes the input labels and returns updated classes and texts. |
Examples:
>>> loader = RandomLoadText(prompt_format="Object: {}", neg_samples=(5, 10), max_samples=20)
>>> labels = {"cls": [0, 1, 2], "texts": [["cat"], ["dog"], ["bird"]], "instances": [...]}
>>> updated_labels = loader(labels)
>>> print(updated_labels["texts"])
['Object: cat', 'Object: dog', 'Object: bird', 'Object: elephant', 'Object: car']
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
prompt_format
|
str
|
Format string for the prompt. Default is '{}'. The format string should contain a single pair of curly braces {} where the text will be inserted. |
'{}'
|
neg_samples
|
Tuple[int, int]
|
A range to randomly sample negative texts. The first integer specifies the minimum number of negative samples, and the second integer specifies the maximum. Default is (80, 80). |
(80, 80)
|
max_samples
|
int
|
The maximum number of different text samples in one image. Default is 80. |
80
|
padding
|
bool
|
Whether to pad texts to max_samples. If True, the number of texts will always be equal to max_samples. Default is False. |
False
|
padding_value
|
str
|
The padding text to use when padding is True. Default is an empty string. |
''
|
Attributes:
Name | Type | Description |
---|---|---|
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
Source code in ultralytics/data/augment.py
__call__
Randomly samples positive and negative texts and updates class indices accordingly.
This method samples positive texts based on the existing class labels in the image, and randomly selects negative texts from the remaining classes. It then updates the class indices to match the new sampled text order.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
labels
|
Dict
|
A dictionary containing image labels and metadata. Must include 'texts' and 'cls' keys. |
required |
Returns:
Type | Description |
---|---|
Dict
|
Updated labels dictionary with new 'cls' and 'texts' entries. |
Examples:
>>> loader = RandomLoadText(prompt_format="A photo of {}", neg_samples=(5, 10), max_samples=20)
>>> labels = {"cls": np.array([[0], [1], [2]]), "texts": [["dog"], ["cat"], ["bird"]]}
>>> updated_labels = loader(labels)
Source code in ultralytics/data/augment.py
2209 2210 2211 2212 2213 2214 2215 2216 2217 2218 2219 2220 2221 2222 2223 2224 2225 2226 2227 2228 2229 2230 2231 2232 2233 2234 2235 2236 2237 2238 2239 2240 2241 2242 2243 2244 2245 2246 2247 2248 2249 2250 2251 2252 2253 2254 2255 2256 2257 2258 2259 2260 2261 2262 2263 2264 2265 2266 2267 2268 2269 2270 |
|
ultralytics.data.augment.ClassifyLetterBox
A class for resizing and padding images for classification tasks.
This class is designed to be part of a transformation pipeline, e.g., T.Compose([LetterBox(size), ToTensor()]). It resizes and pads images to a specified size while maintaining the original aspect ratio.
Attributes:
Name | Type | Description |
---|---|---|
h |
int
|
Target height of the image. |
w |
int
|
Target width of the image. |
auto |
bool
|
If True, automatically calculates the short side using stride. |
stride |
int
|
The stride value, used when 'auto' is True. |
Methods:
Name | Description |
---|---|
__call__ |
Applies the letterbox transformation to an input image. |
Examples:
>>> transform = ClassifyLetterBox(size=(640, 640), auto=False, stride=32)
>>> img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
>>> result = transform(img)
>>> print(result.shape)
(640, 640, 3)
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
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). |
(640, 640)
|
auto
|
bool
|
If True, automatically calculates the short side based on stride. Default is False. |
False
|
stride
|
int
|
The stride value, used when 'auto' is True. Default is 32. |
32
|
Attributes:
Name | Type | Description |
---|---|---|
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)
Source code in ultralytics/data/augment.py
__call__
Resizes and pads an image using the letterbox method.
This method resizes the input image to fit within the specified dimensions while maintaining its aspect ratio, then pads the resized image to match the target size.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
im
|
ndarray
|
Input image as a numpy array with shape (H, W, C). |
required |
Returns:
Type | Description |
---|---|
ndarray
|
Resized and padded image as a numpy array with shape (hs, ws, 3), where hs and ws are the target height and width respectively. |
Examples:
>>> letterbox = ClassifyLetterBox(size=(640, 640))
>>> image = np.random.randint(0, 255, (720, 1280, 3), dtype=np.uint8)
>>> resized_image = letterbox(image)
>>> print(resized_image.shape)
(640, 640, 3)
Source code in ultralytics/data/augment.py
ultralytics.data.augment.CenterCrop
Applies center cropping to images for classification tasks.
This class performs center cropping on input images, resizing them to a specified size while maintaining the aspect ratio. It is designed to be part of a transformation pipeline, e.g., T.Compose([CenterCrop(size), ToTensor()]).
Attributes:
Name | Type | Description |
---|---|---|
h |
int
|
Target height of the cropped image. |
w |
int
|
Target width of the cropped image. |
Methods:
Name | Description |
---|---|
__call__ |
Applies the center crop transformation to an input image. |
Examples:
>>> transform = CenterCrop(640)
>>> image = np.random.randint(0, 255, (1080, 1920, 3), dtype=np.uint8)
>>> cropped_image = transform(image)
>>> print(cropped_image.shape)
(640, 640, 3)
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
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. |
640
|
Returns:
Type | Description |
---|---|
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)
Source code in ultralytics/data/augment.py
__call__
Applies center cropping to an input image.
This method resizes and crops the center of the image using a letterbox method. It maintains the aspect ratio of the original image while fitting it into the specified dimensions.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
im
|
ndarray | Image
|
The input image as a numpy array of shape (H, W, C) or a PIL Image object. |
required |
Returns:
Type | Description |
---|---|
ndarray
|
The center-cropped and resized image as a numpy array of shape (self.h, self.w, C). |
Examples:
>>> transform = CenterCrop(size=224)
>>> image = np.random.randint(0, 255, (640, 480, 3), dtype=np.uint8)
>>> cropped_image = transform(image)
>>> assert cropped_image.shape == (224, 224, 3)
Source code in ultralytics/data/augment.py
ultralytics.data.augment.ToTensor
Converts an image from a numpy array to a PyTorch tensor.
This class is designed to be part of a transformation pipeline, e.g., T.Compose([LetterBox(size), ToTensor()]).
Attributes:
Name | Type | Description |
---|---|---|
half |
bool
|
If True, converts the image to half precision (float16). |
Methods:
Name | Description |
---|---|
__call__ |
Applies the tensor conversion to an input image. |
Examples:
>>> transform = ToTensor(half=True)
>>> img = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
>>> tensor_img = transform(img)
>>> print(tensor_img.shape, tensor_img.dtype)
torch.Size([3, 640, 640]) torch.float16
Notes
The input image is expected to be in BGR format with shape (H, W, C). The output tensor will be in RGB format with shape (C, H, W), normalized to [0, 1].
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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
half
|
bool
|
If True, converts the tensor to half precision (float16). Default is False. |
False
|
Examples:
>>> transform = ToTensor(half=True)
>>> img = np.random.rand(640, 640, 3)
>>> tensor_img = transform(img)
>>> print(tensor_img.dtype)
torch.float16
Source code in ultralytics/data/augment.py
__call__
Transforms an image from a numpy array to a PyTorch tensor.
This method converts the input image from a numpy array to a PyTorch tensor, applying optional half-precision conversion and normalization. The image is transposed from HWC to CHW format and the color channels are reversed from BGR to RGB.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
im
|
ndarray
|
Input image as a numpy array with shape (H, W, C) in BGR order. |
required |
Returns:
Type | Description |
---|---|
Tensor
|
The transformed image as a PyTorch tensor in float32 or float16, normalized to [0, 1] with shape (C, H, W) in RGB order. |
Examples:
>>> transform = ToTensor(half=True)
>>> img = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
>>> tensor_img = transform(img)
>>> print(tensor_img.shape, tensor_img.dtype)
torch.Size([3, 640, 640]) torch.float16
Source code in ultralytics/data/augment.py
ultralytics.data.augment.v8_transforms
Applies a series of image transformations for training.
This function creates a composition of image augmentation techniques to prepare images for YOLO training. It includes operations such as mosaic, copy-paste, random perspective, mixup, and various color adjustments.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset
|
Dataset
|
The dataset object containing image data and annotations. |
required |
imgsz
|
int
|
The target image size for resizing. |
required |
hyp
|
Dict
|
A dictionary of hyperparameters controlling various aspects of the transformations. |
required |
stretch
|
bool
|
If True, applies stretching to the image. If False, uses LetterBox resizing. |
False
|
Returns:
Type | Description |
---|---|
Compose
|
A composition of image transformations to be applied to the dataset. |
Examples:
>>> from ultralytics.data.dataset import YOLODataset
>>> dataset = YOLODataset(img_path="path/to/images", imgsz=640)
>>> hyp = {"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])
Source code in ultralytics/data/augment.py
2273 2274 2275 2276 2277 2278 2279 2280 2281 2282 2283 2284 2285 2286 2287 2288 2289 2290 2291 2292 2293 2294 2295 2296 2297 2298 2299 2300 2301 2302 2303 2304 2305 2306 2307 2308 2309 2310 2311 2312 2313 2314 2315 2316 2317 2318 2319 2320 2321 2322 2323 2324 2325 2326 2327 2328 2329 2330 2331 2332 2333 2334 2335 2336 |
|
ultralytics.data.augment.classify_transforms
classify_transforms(size=224, mean=DEFAULT_MEAN, std=DEFAULT_STD, interpolation='BILINEAR', crop_fraction: float = DEFAULT_CROP_FRACTION)
Creates a composition of image transforms for classification tasks.
This function generates a sequence of torchvision transforms suitable for preprocessing images for classification models during evaluation or inference. The transforms include resizing, center cropping, conversion to tensor, and normalization.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
size
|
int | tuple
|
The target size for the transformed image. If an int, it defines the shortest edge. If a tuple, it defines (height, width). |
224
|
mean
|
tuple
|
Mean values for each RGB channel used in normalization. |
DEFAULT_MEAN
|
std
|
tuple
|
Standard deviation values for each RGB channel used in normalization. |
DEFAULT_STD
|
interpolation
|
str
|
Interpolation method of either 'NEAREST', 'BILINEAR' or 'BICUBIC'. |
'BILINEAR'
|
crop_fraction
|
float
|
Fraction of the image to be cropped. |
DEFAULT_CROP_FRACTION
|
Returns:
Type | Description |
---|---|
Compose
|
A composition of torchvision transforms. |
Examples:
>>> transforms = classify_transforms(size=224)
>>> img = Image.open("path/to/image.jpg")
>>> transformed_img = transforms(img)
Source code in ultralytics/data/augment.py
ultralytics.data.augment.classify_augmentations
classify_augmentations(size=224, mean=DEFAULT_MEAN, std=DEFAULT_STD, scale=None, ratio=None, hflip=0.5, vflip=0.0, auto_augment=None, hsv_h=0.015, hsv_s=0.4, hsv_v=0.4, force_color_jitter=False, erasing=0.0, interpolation='BILINEAR')
Creates a composition of image augmentation transforms for classification tasks.
This function generates a set of image transformations suitable for training classification models. It includes options for resizing, flipping, color jittering, auto augmentation, and random erasing.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
size
|
int
|
Target size for the image after transformations. |
224
|
mean
|
tuple
|
Mean values for normalization, one per channel. |
DEFAULT_MEAN
|
std
|
tuple
|
Standard deviation values for normalization, one per channel. |
DEFAULT_STD
|
scale
|
tuple | None
|
Range of size of the origin size cropped. |
None
|
ratio
|
tuple | None
|
Range of aspect ratio of the origin aspect ratio cropped. |
None
|
hflip
|
float
|
Probability of horizontal flip. |
0.5
|
vflip
|
float
|
Probability of vertical flip. |
0.0
|
auto_augment
|
str | None
|
Auto augmentation policy. Can be 'randaugment', 'augmix', 'autoaugment' or None. |
None
|
hsv_h
|
float
|
Image HSV-Hue augmentation factor. |
0.015
|
hsv_s
|
float
|
Image HSV-Saturation augmentation factor. |
0.4
|
hsv_v
|
float
|
Image HSV-Value augmentation factor. |
0.4
|
force_color_jitter
|
bool
|
Whether to apply color jitter even if auto augment is enabled. |
False
|
erasing
|
float
|
Probability of random erasing. |
0.0
|
interpolation
|
str
|
Interpolation method of either 'NEAREST', 'BILINEAR' or 'BICUBIC'. |
'BILINEAR'
|
Returns:
Type | Description |
---|---|
Compose
|
A composition of image augmentation transforms. |
Examples:
>>> transforms = classify_augmentations(size=224, auto_augment="randaugment")
>>> augmented_image = transforms(original_image)
Source code in ultralytics/data/augment.py
2397 2398 2399 2400 2401 2402 2403 2404 2405 2406 2407 2408 2409 2410 2411 2412 2413 2414 2415 2416 2417 2418 2419 2420 2421 2422 2423 2424 2425 2426 2427 2428 2429 2430 2431 2432 2433 2434 2435 2436 2437 2438 2439 2440 2441 2442 2443 2444 2445 2446 2447 2448 2449 2450 2451 2452 2453 2454 2455 2456 2457 2458 2459 2460 2461 2462 2463 2464 2465 2466 2467 2468 2469 2470 2471 2472 2473 2474 2475 2476 2477 2478 2479 2480 2481 2482 2483 2484 2485 2486 2487 2488 2489 2490 2491 2492 2493 2494 2495 2496 2497 |
|