Reference for ultralytics/data/base.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/data/base.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!
ultralytics.data.base.BaseDataset
BaseDataset(
img_path,
imgsz=640,
cache=False,
augment=True,
hyp=DEFAULT_CFG,
prefix="",
rect=False,
batch_size=16,
stride=32,
pad=0.5,
single_cls=False,
classes=None,
fraction=1.0,
)
Bases: Dataset
Base dataset class for loading and processing image data.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
img_path
|
str
|
Path to the folder containing images. |
required |
imgsz
|
int
|
Image size. Defaults to 640. |
640
|
cache
|
bool
|
Cache images to RAM or disk during training. Defaults to False. |
False
|
augment
|
bool
|
If True, data augmentation is applied. Defaults to True. |
True
|
hyp
|
dict
|
Hyperparameters to apply data augmentation. Defaults to None. |
DEFAULT_CFG
|
prefix
|
str
|
Prefix to print in log messages. Defaults to ''. |
''
|
rect
|
bool
|
If True, rectangular training is used. Defaults to False. |
False
|
batch_size
|
int
|
Size of batches. Defaults to None. |
16
|
stride
|
int
|
Stride. Defaults to 32. |
32
|
pad
|
float
|
Padding. Defaults to 0.0. |
0.5
|
single_cls
|
bool
|
If True, single class training is used. Defaults to False. |
False
|
classes
|
list
|
List of included classes. Default is None. |
None
|
fraction
|
float
|
Fraction of dataset to utilize. Default is 1.0 (use all data). |
1.0
|
Attributes:
Name | Type | Description |
---|---|---|
im_files |
list
|
List of image file paths. |
labels |
list
|
List of label data dictionaries. |
ni |
int
|
Number of images in the dataset. |
ims |
list
|
List of loaded images. |
npy_files |
list
|
List of numpy file paths. |
transforms |
callable
|
Image transformation function. |
Source code in ultralytics/data/base.py
__getitem__
__len__
build_transforms
Users can customize augmentations here.
Example
Source code in ultralytics/data/base.py
cache_images
Cache images to memory or disk.
Source code in ultralytics/data/base.py
cache_images_to_disk
Saves an image as an *.npy file for faster loading.
check_cache_disk
Check image caching requirements vs available disk space.
Source code in ultralytics/data/base.py
check_cache_ram
Check image caching requirements vs available memory.
Source code in ultralytics/data/base.py
get_image_and_label
Get and return label information from the dataset.
Source code in ultralytics/data/base.py
get_img_files
Read image files.
Source code in ultralytics/data/base.py
get_labels
Users can customize their own format here.
Note
Ensure output is a dictionary with the following keys:
Source code in ultralytics/data/base.py
load_image
Loads 1 image from dataset index 'i', returns (im, resized hw).
Source code in ultralytics/data/base.py
set_rectangle
Sets the shape of bounding boxes for YOLO detections as rectangles.
Source code in ultralytics/data/base.py
update_labels
Update labels to include only these classes (optional).