Skip to content

Reference for ultralytics/data/dataset.py

Note

Full source code for this file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/data/dataset.py. Help us fix any issues you see by submitting a Pull Request 🛠️. Thank you 🙏!


ultralytics.data.dataset.YOLODataset

Bases: BaseDataset

Dataset class for loading object detection and/or segmentation labels in YOLO format.

Parameters:

Name Type Description Default
data dict

A dataset YAML dictionary. Defaults to None.

None
use_segments bool

If True, segmentation masks are used as labels. Defaults to False.

False
use_keypoints bool

If True, keypoints are used as labels. Defaults to False.

False

Returns:

Type Description
Dataset

A PyTorch dataset object that can be used for training an object detection model.

Source code in ultralytics/data/dataset.py
class YOLODataset(BaseDataset):
    """
    Dataset class for loading object detection and/or segmentation labels in YOLO format.

    Args:
        data (dict, optional): A dataset YAML dictionary. Defaults to None.
        use_segments (bool, optional): If True, segmentation masks are used as labels. Defaults to False.
        use_keypoints (bool, optional): If True, keypoints are used as labels. Defaults to False.

    Returns:
        (torch.utils.data.Dataset): A PyTorch dataset object that can be used for training an object detection model.
    """

    def __init__(self, *args, data=None, use_segments=False, use_keypoints=False, **kwargs):
        self.use_segments = use_segments
        self.use_keypoints = use_keypoints
        self.data = data
        assert not (self.use_segments and self.use_keypoints), 'Can not use both segments and keypoints.'
        super().__init__(*args, **kwargs)

    def cache_labels(self, path=Path('./labels.cache')):
        """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}...'
        total = len(self.im_files)
        nkpt, ndim = self.data.get('kpt_shape', (0, 0))
        if self.use_keypoints and (nkpt <= 0 or ndim not in (2, 3)):
            raise ValueError("'kpt_shape' in data.yaml missing or incorrect. Should be a list with [number of "
                             "keypoints, number of dims (2 for x,y or 3 for x,y,visible)], i.e. 'kpt_shape: [17, 3]'")
        with ThreadPool(NUM_THREADS) as pool:
            results = pool.imap(func=verify_image_label,
                                iterable=zip(self.im_files, self.label_files, repeat(self.prefix),
                                             repeat(self.use_keypoints), repeat(len(self.data['names'])), repeat(nkpt),
                                             repeat(ndim)))
            pbar = TQDM(results, desc=desc, total=total)
            for im_file, lb, shape, segments, keypoint, nm_f, nf_f, ne_f, nc_f, msg in pbar:
                nm += nm_f
                nf += nf_f
                ne += ne_f
                nc += nc_f
                if im_file:
                    x['labels'].append(
                        dict(
                            im_file=im_file,
                            shape=shape,
                            cls=lb[:, 0:1],  # n, 1
                            bboxes=lb[:, 1:],  # n, 4
                            segments=segments,
                            keypoints=keypoint,
                            normalized=True,
                            bbox_format='xywh'))
                if msg:
                    msgs.append(msg)
                pbar.desc = f'{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt'
            pbar.close()

        if msgs:
            LOGGER.info('\n'.join(msgs))
        if nf == 0:
            LOGGER.warning(f'{self.prefix}WARNING ⚠️ No labels found in {path}. {HELP_URL}')
        x['hash'] = get_hash(self.label_files + self.im_files)
        x['results'] = nf, nm, ne, nc, len(self.im_files)
        x['msgs'] = msgs  # warnings
        save_dataset_cache_file(self.prefix, path, x)
        return x

    def get_labels(self):
        """Returns dictionary of labels for YOLO training."""
        self.label_files = img2label_paths(self.im_files)
        cache_path = Path(self.label_files[0]).parent.with_suffix('.cache')
        try:
            cache, exists = load_dataset_cache_file(cache_path), True  # attempt to load a *.cache file
            assert cache['version'] == DATASET_CACHE_VERSION  # matches current version
            assert cache['hash'] == get_hash(self.label_files + self.im_files)  # identical hash
        except (FileNotFoundError, AssertionError, AttributeError):
            cache, exists = self.cache_labels(cache_path), False  # run cache ops

        # Display cache
        nf, nm, ne, nc, n = cache.pop('results')  # found, missing, empty, corrupt, total
        if exists and LOCAL_RANK in (-1, 0):
            d = f'Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt'
            TQDM(None, desc=self.prefix + d, total=n, initial=n)  # display results
            if cache['msgs']:
                LOGGER.info('\n'.join(cache['msgs']))  # display warnings

        # Read cache
        [cache.pop(k) for k in ('hash', 'version', 'msgs')]  # remove items
        labels = cache['labels']
        if not labels:
            LOGGER.warning(f'WARNING ⚠️ No images found in {cache_path}, training may not work correctly. {HELP_URL}')
        self.im_files = [lb['im_file'] for lb in labels]  # update im_files

        # Check if the dataset is all boxes or all segments
        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}, '
                f'len(boxes) = {len_boxes}. To resolve this only boxes will be used and all segments will be removed. '
                'To avoid this please supply either a detect or segment dataset, not a detect-segment mixed dataset.')
            for lb in labels:
                lb['segments'] = []
        if len_cls == 0:
            LOGGER.warning(f'WARNING ⚠️ No labels found in {cache_path}, training may not work correctly. {HELP_URL}')
        return labels

    def build_transforms(self, hyp=None):
        """Builds and appends transforms to the list."""
        if self.augment:
            hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0
            hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0
            transforms = v8_transforms(self, self.imgsz, hyp)
        else:
            transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), scaleup=False)])
        transforms.append(
            Format(bbox_format='xywh',
                   normalize=True,
                   return_mask=self.use_segments,
                   return_keypoint=self.use_keypoints,
                   batch_idx=True,
                   mask_ratio=hyp.mask_ratio,
                   mask_overlap=hyp.overlap_mask))
        return transforms

    def close_mosaic(self, hyp):
        """Sets mosaic, copy_paste and mixup options to 0.0 and builds transformations."""
        hyp.mosaic = 0.0  # set mosaic ratio=0.0
        hyp.copy_paste = 0.0  # keep the same behavior as previous v8 close-mosaic
        hyp.mixup = 0.0  # keep the same behavior as previous v8 close-mosaic
        self.transforms = self.build_transforms(hyp)

    def update_labels_info(self, label):
        """custom your label format here."""
        # NOTE: cls is not with bboxes now, classification and semantic segmentation need an independent cls label
        # we can make it also support classification and semantic segmentation by add or remove some dict keys there.
        bboxes = label.pop('bboxes')
        segments = label.pop('segments')
        keypoints = label.pop('keypoints', None)
        bbox_format = label.pop('bbox_format')
        normalized = label.pop('normalized')
        label['instances'] = Instances(bboxes, segments, keypoints, bbox_format=bbox_format, normalized=normalized)
        return label

    @staticmethod
    def collate_fn(batch):
        """Collates data samples into batches."""
        new_batch = {}
        keys = batch[0].keys()
        values = list(zip(*[list(b.values()) for b in batch]))
        for i, k in enumerate(keys):
            value = values[i]
            if k == 'img':
                value = torch.stack(value, 0)
            if k in ['masks', 'keypoints', 'bboxes', 'cls']:
                value = torch.cat(value, 0)
            new_batch[k] = value
        new_batch['batch_idx'] = list(new_batch['batch_idx'])
        for i in range(len(new_batch['batch_idx'])):
            new_batch['batch_idx'][i] += i  # add target image index for build_targets()
        new_batch['batch_idx'] = torch.cat(new_batch['batch_idx'], 0)
        return new_batch

build_transforms(hyp=None)

Builds and appends transforms to the list.

Source code in ultralytics/data/dataset.py
def build_transforms(self, hyp=None):
    """Builds and appends transforms to the list."""
    if self.augment:
        hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0
        hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0
        transforms = v8_transforms(self, self.imgsz, hyp)
    else:
        transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), scaleup=False)])
    transforms.append(
        Format(bbox_format='xywh',
               normalize=True,
               return_mask=self.use_segments,
               return_keypoint=self.use_keypoints,
               batch_idx=True,
               mask_ratio=hyp.mask_ratio,
               mask_overlap=hyp.overlap_mask))
    return transforms

cache_labels(path=Path('./labels.cache'))

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.

Source code in ultralytics/data/dataset.py
def cache_labels(self, path=Path('./labels.cache')):
    """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}...'
    total = len(self.im_files)
    nkpt, ndim = self.data.get('kpt_shape', (0, 0))
    if self.use_keypoints and (nkpt <= 0 or ndim not in (2, 3)):
        raise ValueError("'kpt_shape' in data.yaml missing or incorrect. Should be a list with [number of "
                         "keypoints, number of dims (2 for x,y or 3 for x,y,visible)], i.e. 'kpt_shape: [17, 3]'")
    with ThreadPool(NUM_THREADS) as pool:
        results = pool.imap(func=verify_image_label,
                            iterable=zip(self.im_files, self.label_files, repeat(self.prefix),
                                         repeat(self.use_keypoints), repeat(len(self.data['names'])), repeat(nkpt),
                                         repeat(ndim)))
        pbar = TQDM(results, desc=desc, total=total)
        for im_file, lb, shape, segments, keypoint, nm_f, nf_f, ne_f, nc_f, msg in pbar:
            nm += nm_f
            nf += nf_f
            ne += ne_f
            nc += nc_f
            if im_file:
                x['labels'].append(
                    dict(
                        im_file=im_file,
                        shape=shape,
                        cls=lb[:, 0:1],  # n, 1
                        bboxes=lb[:, 1:],  # n, 4
                        segments=segments,
                        keypoints=keypoint,
                        normalized=True,
                        bbox_format='xywh'))
            if msg:
                msgs.append(msg)
            pbar.desc = f'{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt'
        pbar.close()

    if msgs:
        LOGGER.info('\n'.join(msgs))
    if nf == 0:
        LOGGER.warning(f'{self.prefix}WARNING ⚠️ No labels found in {path}. {HELP_URL}')
    x['hash'] = get_hash(self.label_files + self.im_files)
    x['results'] = nf, nm, ne, nc, len(self.im_files)
    x['msgs'] = msgs  # warnings
    save_dataset_cache_file(self.prefix, path, x)
    return x

close_mosaic(hyp)

Sets mosaic, copy_paste and mixup options to 0.0 and builds transformations.

Source code in ultralytics/data/dataset.py
def close_mosaic(self, hyp):
    """Sets mosaic, copy_paste and mixup options to 0.0 and builds transformations."""
    hyp.mosaic = 0.0  # set mosaic ratio=0.0
    hyp.copy_paste = 0.0  # keep the same behavior as previous v8 close-mosaic
    hyp.mixup = 0.0  # keep the same behavior as previous v8 close-mosaic
    self.transforms = self.build_transforms(hyp)

collate_fn(batch) staticmethod

Collates data samples into batches.

Source code in ultralytics/data/dataset.py
@staticmethod
def collate_fn(batch):
    """Collates data samples into batches."""
    new_batch = {}
    keys = batch[0].keys()
    values = list(zip(*[list(b.values()) for b in batch]))
    for i, k in enumerate(keys):
        value = values[i]
        if k == 'img':
            value = torch.stack(value, 0)
        if k in ['masks', 'keypoints', 'bboxes', 'cls']:
            value = torch.cat(value, 0)
        new_batch[k] = value
    new_batch['batch_idx'] = list(new_batch['batch_idx'])
    for i in range(len(new_batch['batch_idx'])):
        new_batch['batch_idx'][i] += i  # add target image index for build_targets()
    new_batch['batch_idx'] = torch.cat(new_batch['batch_idx'], 0)
    return new_batch

get_labels()

Returns dictionary of labels for YOLO training.

Source code in ultralytics/data/dataset.py
def get_labels(self):
    """Returns dictionary of labels for YOLO training."""
    self.label_files = img2label_paths(self.im_files)
    cache_path = Path(self.label_files[0]).parent.with_suffix('.cache')
    try:
        cache, exists = load_dataset_cache_file(cache_path), True  # attempt to load a *.cache file
        assert cache['version'] == DATASET_CACHE_VERSION  # matches current version
        assert cache['hash'] == get_hash(self.label_files + self.im_files)  # identical hash
    except (FileNotFoundError, AssertionError, AttributeError):
        cache, exists = self.cache_labels(cache_path), False  # run cache ops

    # Display cache
    nf, nm, ne, nc, n = cache.pop('results')  # found, missing, empty, corrupt, total
    if exists and LOCAL_RANK in (-1, 0):
        d = f'Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt'
        TQDM(None, desc=self.prefix + d, total=n, initial=n)  # display results
        if cache['msgs']:
            LOGGER.info('\n'.join(cache['msgs']))  # display warnings

    # Read cache
    [cache.pop(k) for k in ('hash', 'version', 'msgs')]  # remove items
    labels = cache['labels']
    if not labels:
        LOGGER.warning(f'WARNING ⚠️ No images found in {cache_path}, training may not work correctly. {HELP_URL}')
    self.im_files = [lb['im_file'] for lb in labels]  # update im_files

    # Check if the dataset is all boxes or all segments
    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}, '
            f'len(boxes) = {len_boxes}. To resolve this only boxes will be used and all segments will be removed. '
            'To avoid this please supply either a detect or segment dataset, not a detect-segment mixed dataset.')
        for lb in labels:
            lb['segments'] = []
    if len_cls == 0:
        LOGGER.warning(f'WARNING ⚠️ No labels found in {cache_path}, training may not work correctly. {HELP_URL}')
    return labels

update_labels_info(label)

custom your label format here.

Source code in ultralytics/data/dataset.py
def update_labels_info(self, label):
    """custom your label format here."""
    # NOTE: cls is not with bboxes now, classification and semantic segmentation need an independent cls label
    # we can make it also support classification and semantic segmentation by add or remove some dict keys there.
    bboxes = label.pop('bboxes')
    segments = label.pop('segments')
    keypoints = label.pop('keypoints', None)
    bbox_format = label.pop('bbox_format')
    normalized = label.pop('normalized')
    label['instances'] = Instances(bboxes, segments, keypoints, bbox_format=bbox_format, normalized=normalized)
    return label




ultralytics.data.dataset.ClassificationDataset

Bases: ImageFolder

YOLO Classification Dataset.

Parameters:

Name Type Description Default
root str

Dataset path.

required

Attributes:

Name Type Description
cache_ram bool

True if images should be cached in RAM, False otherwise.

cache_disk bool

True if images should be cached on disk, False otherwise.

samples list

List of samples containing file, index, npy, and im.

torch_transforms callable

torchvision transforms applied to the dataset.

album_transforms callable

Albumentations transforms applied to the dataset if augment is True.

Source code in ultralytics/data/dataset.py
class ClassificationDataset(torchvision.datasets.ImageFolder):
    """
    YOLO Classification Dataset.

    Args:
        root (str): Dataset path.

    Attributes:
        cache_ram (bool): True if images should be cached in RAM, False otherwise.
        cache_disk (bool): True if images should be cached on disk, False otherwise.
        samples (list): List of samples containing file, index, npy, and im.
        torch_transforms (callable): torchvision transforms applied to the dataset.
        album_transforms (callable, optional): Albumentations transforms applied to the dataset if augment is True.
    """

    def __init__(self, root, args, augment=False, cache=False, prefix=''):
        """
        Initialize YOLO object with root, image size, augmentations, and cache settings.

        Args:
            root (str): Dataset path.
            args (Namespace): Argument parser containing dataset related settings.
            augment (bool, optional): True if dataset should be augmented, False otherwise. Defaults to False.
            cache (bool | str | optional): Cache setting, can be True, False, 'ram' or 'disk'. Defaults to False.
        """
        super().__init__(root=root)
        if augment and args.fraction < 1.0:  # reduce training fraction
            self.samples = self.samples[:round(len(self.samples) * args.fraction)]
        self.prefix = colorstr(f'{prefix}: ') if prefix else ''
        self.cache_ram = cache is True or cache == 'ram'
        self.cache_disk = cache == 'disk'
        self.samples = self.verify_images()  # filter out bad images
        self.samples = [list(x) + [Path(x[0]).with_suffix('.npy'), None] for x in self.samples]  # file, index, npy, im
        self.torch_transforms = classify_transforms(args.imgsz)
        self.album_transforms = classify_albumentations(
            augment=augment,
            size=args.imgsz,
            scale=(1.0 - args.scale, 1.0),  # (0.08, 1.0)
            hflip=args.fliplr,
            vflip=args.flipud,
            hsv_h=args.hsv_h,  # HSV-Hue augmentation (fraction)
            hsv_s=args.hsv_s,  # HSV-Saturation augmentation (fraction)
            hsv_v=args.hsv_v,  # HSV-Value augmentation (fraction)
            mean=(0.0, 0.0, 0.0),  # IMAGENET_MEAN
            std=(1.0, 1.0, 1.0),  # IMAGENET_STD
            auto_aug=False) if augment else None

    def __getitem__(self, i):
        """Returns subset of data and targets corresponding to given indices."""
        f, j, fn, im = self.samples[i]  # filename, index, filename.with_suffix('.npy'), image
        if self.cache_ram and im is None:
            im = self.samples[i][3] = cv2.imread(f)
        elif self.cache_disk:
            if not fn.exists():  # load npy
                np.save(fn.as_posix(), cv2.imread(f), allow_pickle=False)
            im = np.load(fn)
        else:  # read image
            im = cv2.imread(f)  # BGR
        if self.album_transforms:
            sample = self.album_transforms(image=cv2.cvtColor(im, cv2.COLOR_BGR2RGB))['image']
        else:
            sample = self.torch_transforms(im)
        return {'img': sample, 'cls': j}

    def __len__(self) -> int:
        return len(self.samples)

    def verify_images(self):
        """Verify all images in dataset."""
        desc = f'{self.prefix}Scanning {self.root}...'
        path = Path(self.root).with_suffix('.cache')  # *.cache file path

        with contextlib.suppress(FileNotFoundError, AssertionError, AttributeError):
            cache = load_dataset_cache_file(path)  # attempt to load a *.cache file
            assert cache['version'] == DATASET_CACHE_VERSION  # matches current version
            assert cache['hash'] == get_hash([x[0] for x in self.samples])  # identical hash
            nf, nc, n, samples = cache.pop('results')  # found, missing, empty, corrupt, total
            if LOCAL_RANK in (-1, 0):
                d = f'{desc} {nf} images, {nc} corrupt'
                TQDM(None, desc=d, total=n, initial=n)
                if cache['msgs']:
                    LOGGER.info('\n'.join(cache['msgs']))  # display warnings
            return samples

        # Run scan if *.cache retrieval failed
        nf, nc, msgs, samples, x = 0, 0, [], [], {}
        with ThreadPool(NUM_THREADS) as pool:
            results = pool.imap(func=verify_image, iterable=zip(self.samples, repeat(self.prefix)))
            pbar = TQDM(results, desc=desc, total=len(self.samples))
            for sample, nf_f, nc_f, msg in pbar:
                if nf_f:
                    samples.append(sample)
                if msg:
                    msgs.append(msg)
                nf += nf_f
                nc += nc_f
                pbar.desc = f'{desc} {nf} images, {nc} corrupt'
            pbar.close()
        if msgs:
            LOGGER.info('\n'.join(msgs))
        x['hash'] = get_hash([x[0] for x in self.samples])
        x['results'] = nf, nc, len(samples), samples
        x['msgs'] = msgs  # warnings
        save_dataset_cache_file(self.prefix, path, x)
        return samples

__getitem__(i)

Returns subset of data and targets corresponding to given indices.

Source code in ultralytics/data/dataset.py
def __getitem__(self, i):
    """Returns subset of data and targets corresponding to given indices."""
    f, j, fn, im = self.samples[i]  # filename, index, filename.with_suffix('.npy'), image
    if self.cache_ram and im is None:
        im = self.samples[i][3] = cv2.imread(f)
    elif self.cache_disk:
        if not fn.exists():  # load npy
            np.save(fn.as_posix(), cv2.imread(f), allow_pickle=False)
        im = np.load(fn)
    else:  # read image
        im = cv2.imread(f)  # BGR
    if self.album_transforms:
        sample = self.album_transforms(image=cv2.cvtColor(im, cv2.COLOR_BGR2RGB))['image']
    else:
        sample = self.torch_transforms(im)
    return {'img': sample, 'cls': j}

__init__(root, args, augment=False, cache=False, prefix='')

Initialize YOLO object with root, image size, augmentations, and cache settings.

Parameters:

Name Type Description Default
root str

Dataset path.

required
args Namespace

Argument parser containing dataset related settings.

required
augment bool

True if dataset should be augmented, False otherwise. Defaults to False.

False
cache bool | str | optional

Cache setting, can be True, False, 'ram' or 'disk'. Defaults to False.

False
Source code in ultralytics/data/dataset.py
def __init__(self, root, args, augment=False, cache=False, prefix=''):
    """
    Initialize YOLO object with root, image size, augmentations, and cache settings.

    Args:
        root (str): Dataset path.
        args (Namespace): Argument parser containing dataset related settings.
        augment (bool, optional): True if dataset should be augmented, False otherwise. Defaults to False.
        cache (bool | str | optional): Cache setting, can be True, False, 'ram' or 'disk'. Defaults to False.
    """
    super().__init__(root=root)
    if augment and args.fraction < 1.0:  # reduce training fraction
        self.samples = self.samples[:round(len(self.samples) * args.fraction)]
    self.prefix = colorstr(f'{prefix}: ') if prefix else ''
    self.cache_ram = cache is True or cache == 'ram'
    self.cache_disk = cache == 'disk'
    self.samples = self.verify_images()  # filter out bad images
    self.samples = [list(x) + [Path(x[0]).with_suffix('.npy'), None] for x in self.samples]  # file, index, npy, im
    self.torch_transforms = classify_transforms(args.imgsz)
    self.album_transforms = classify_albumentations(
        augment=augment,
        size=args.imgsz,
        scale=(1.0 - args.scale, 1.0),  # (0.08, 1.0)
        hflip=args.fliplr,
        vflip=args.flipud,
        hsv_h=args.hsv_h,  # HSV-Hue augmentation (fraction)
        hsv_s=args.hsv_s,  # HSV-Saturation augmentation (fraction)
        hsv_v=args.hsv_v,  # HSV-Value augmentation (fraction)
        mean=(0.0, 0.0, 0.0),  # IMAGENET_MEAN
        std=(1.0, 1.0, 1.0),  # IMAGENET_STD
        auto_aug=False) if augment else None

verify_images()

Verify all images in dataset.

Source code in ultralytics/data/dataset.py
def verify_images(self):
    """Verify all images in dataset."""
    desc = f'{self.prefix}Scanning {self.root}...'
    path = Path(self.root).with_suffix('.cache')  # *.cache file path

    with contextlib.suppress(FileNotFoundError, AssertionError, AttributeError):
        cache = load_dataset_cache_file(path)  # attempt to load a *.cache file
        assert cache['version'] == DATASET_CACHE_VERSION  # matches current version
        assert cache['hash'] == get_hash([x[0] for x in self.samples])  # identical hash
        nf, nc, n, samples = cache.pop('results')  # found, missing, empty, corrupt, total
        if LOCAL_RANK in (-1, 0):
            d = f'{desc} {nf} images, {nc} corrupt'
            TQDM(None, desc=d, total=n, initial=n)
            if cache['msgs']:
                LOGGER.info('\n'.join(cache['msgs']))  # display warnings
        return samples

    # Run scan if *.cache retrieval failed
    nf, nc, msgs, samples, x = 0, 0, [], [], {}
    with ThreadPool(NUM_THREADS) as pool:
        results = pool.imap(func=verify_image, iterable=zip(self.samples, repeat(self.prefix)))
        pbar = TQDM(results, desc=desc, total=len(self.samples))
        for sample, nf_f, nc_f, msg in pbar:
            if nf_f:
                samples.append(sample)
            if msg:
                msgs.append(msg)
            nf += nf_f
            nc += nc_f
            pbar.desc = f'{desc} {nf} images, {nc} corrupt'
        pbar.close()
    if msgs:
        LOGGER.info('\n'.join(msgs))
    x['hash'] = get_hash([x[0] for x in self.samples])
    x['results'] = nf, nc, len(samples), samples
    x['msgs'] = msgs  # warnings
    save_dataset_cache_file(self.prefix, path, x)
    return samples




ultralytics.data.dataset.SemanticDataset

Bases: BaseDataset

Source code in ultralytics/data/dataset.py
class SemanticDataset(BaseDataset):

    def __init__(self):
        """Initialize a SemanticDataset object."""
        super().__init__()

__init__()

Initialize a SemanticDataset object.

Source code in ultralytics/data/dataset.py
def __init__(self):
    """Initialize a SemanticDataset object."""
    super().__init__()




ultralytics.data.dataset.load_dataset_cache_file(path)

Load an Ultralytics *.cache dictionary from path.

Source code in ultralytics/data/dataset.py
def load_dataset_cache_file(path):
    """Load an Ultralytics *.cache dictionary from path."""
    import gc
    gc.disable()  # reduce pickle load time https://github.com/ultralytics/ultralytics/pull/1585
    cache = np.load(str(path), allow_pickle=True).item()  # load dict
    gc.enable()
    return cache




ultralytics.data.dataset.save_dataset_cache_file(prefix, path, x)

Save an Ultralytics dataset *.cache dictionary x to path.

Source code in ultralytics/data/dataset.py
def save_dataset_cache_file(prefix, path, x):
    """Save an Ultralytics dataset *.cache dictionary x to path."""
    x['version'] = DATASET_CACHE_VERSION  # add cache version
    if is_dir_writeable(path.parent):
        if path.exists():
            path.unlink()  # remove *.cache file if exists
        np.save(str(path), x)  # save cache for next time
        path.with_suffix('.cache.npy').rename(path)  # remove .npy suffix
        LOGGER.info(f'{prefix}New cache created: {path}')
    else:
        LOGGER.warning(f'{prefix}WARNING ⚠️ Cache directory {path.parent} is not writeable, cache not saved.')




Created 2023-07-16, Updated 2023-08-24
Authors: glenn-jocher (6), Laughing-q (1)