|
|
|
@ -16,10 +16,28 @@ from .utils import HELP_URL, LOCAL_RANK, get_hash, img2label_paths, verify_image |
|
|
|
|
class YOLODataset(BaseDataset): |
|
|
|
|
cache_version = '1.0.1' # dataset labels *.cache version, >= 1.0.0 for YOLOv8 |
|
|
|
|
rand_interp_methods = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4] |
|
|
|
|
"""YOLO Dataset. |
|
|
|
|
""" |
|
|
|
|
Dataset class for loading images object detection and/or segmentation labels in YOLO format. |
|
|
|
|
|
|
|
|
|
Args: |
|
|
|
|
img_path (str): image path. |
|
|
|
|
prefix (str): prefix. |
|
|
|
|
img_path (str): path to the folder containing images. |
|
|
|
|
imgsz (int): image size (default: 640). |
|
|
|
|
cache (bool): if True, a cache file of the labels is created to speed up future creation of dataset instances |
|
|
|
|
(default: False). |
|
|
|
|
augment (bool): if True, data augmentation is applied (default: True). |
|
|
|
|
hyp (dict): hyperparameters to apply data augmentation (default: None). |
|
|
|
|
prefix (str): prefix to print in log messages (default: ''). |
|
|
|
|
rect (bool): if True, rectangular training is used (default: False). |
|
|
|
|
batch_size (int): size of batches (default: None). |
|
|
|
|
stride (int): stride (default: 32). |
|
|
|
|
pad (float): padding (default: 0.0). |
|
|
|
|
single_cls (bool): if True, single class training is used (default: False). |
|
|
|
|
use_segments (bool): if True, segmentation masks are used as labels (default: False). |
|
|
|
|
use_keypoints (bool): if True, keypoints are used as labels (default: False). |
|
|
|
|
names (list): class names (default: None). |
|
|
|
|
|
|
|
|
|
Returns: |
|
|
|
|
A PyTorch dataset object that can be used for training an object detection or segmentation model. |
|
|
|
|
""" |
|
|
|
|
|
|
|
|
|
def __init__(self, |
|
|
|
@ -44,7 +62,12 @@ class YOLODataset(BaseDataset): |
|
|
|
|
super().__init__(img_path, imgsz, cache, augment, hyp, prefix, rect, batch_size, stride, pad, single_cls) |
|
|
|
|
|
|
|
|
|
def cache_labels(self, path=Path('./labels.cache')): |
|
|
|
|
# Cache dataset labels, check images and read shapes |
|
|
|
|
"""Cache dataset labels, check images and read shapes. |
|
|
|
|
Args: |
|
|
|
|
path (Path): path where to save the cache file (default: Path('./labels.cache')). |
|
|
|
|
Returns: |
|
|
|
|
(dict): labels. |
|
|
|
|
""" |
|
|
|
|
x = {'labels': []} |
|
|
|
|
nm, nf, ne, nc, msgs = 0, 0, 0, 0, [] # number missing, found, empty, corrupt, messages |
|
|
|
|
desc = f'{self.prefix}Scanning {path.parent / path.stem}...' |
|
|
|
@ -119,9 +142,8 @@ class YOLODataset(BaseDataset): |
|
|
|
|
self.im_files = [lb['im_file'] for lb in labels] # update im_files |
|
|
|
|
|
|
|
|
|
# Check if the dataset is all boxes or all segments |
|
|
|
|
len_cls = sum(len(lb['cls']) for lb in labels) |
|
|
|
|
len_boxes = sum(len(lb['bboxes']) for lb in labels) |
|
|
|
|
len_segments = sum(len(lb['segments']) for lb in labels) |
|
|
|
|
lengths = ((len(lb['cls']), len(lb['bboxes']), len(lb['segments'])) for lb in labels) |
|
|
|
|
len_cls, len_boxes, len_segments = (sum(x) for x in zip(*lengths)) |
|
|
|
|
if len_segments and len_boxes != len_segments: |
|
|
|
|
LOGGER.warning( |
|
|
|
|
f'WARNING ⚠️ Box and segment counts should be equal, but got len(segments) = {len_segments}, ' |
|
|
|
|