Skip to content

MixAndRectDataset


A dataset class that applies mosaic and mixup transformations as well as rectangular training.

Attributes:

Name Type Description
dataset

The base dataset.

imgsz

The size of the images in the dataset.

Source code in ultralytics/yolo/data/dataset_wrappers.py
class MixAndRectDataset:
    """
    A dataset class that applies mosaic and mixup transformations as well as rectangular training.

    Attributes:
        dataset: The base dataset.
        imgsz: The size of the images in the dataset.
    """

    def __init__(self, dataset):
        """
        Args:
            dataset (BaseDataset): The base dataset to apply transformations to.
        """
        self.dataset = dataset
        self.imgsz = dataset.imgsz

    def __len__(self):
        """Returns the number of items in the dataset."""
        return len(self.dataset)

    def __getitem__(self, index):
        """
        Applies mosaic, mixup and rectangular training transformations to an item in the dataset.

        Args:
            index (int): Index of the item in the dataset.

        Returns:
            (dict): A dictionary containing the transformed item data.
        """
        labels = deepcopy(self.dataset[index])
        for transform in self.dataset.transforms.tolist():
            # Mosaic and mixup
            if hasattr(transform, 'get_indexes'):
                indexes = transform.get_indexes(self.dataset)
                if not isinstance(indexes, collections.abc.Sequence):
                    indexes = [indexes]
                labels['mix_labels'] = [deepcopy(self.dataset[index]) for index in indexes]
            if self.dataset.rect and isinstance(transform, LetterBox):
                transform.new_shape = self.dataset.batch_shapes[self.dataset.batch[index]]
            labels = transform(labels)
            if 'mix_labels' in labels:
                labels.pop('mix_labels')
        return labels

__getitem__(index)

Applies mosaic, mixup and rectangular training transformations to an item in the dataset.

Parameters:

Name Type Description Default
index int

Index of the item in the dataset.

required

Returns:

Type Description
dict

A dictionary containing the transformed item data.

Source code in ultralytics/yolo/data/dataset_wrappers.py
def __getitem__(self, index):
    """
    Applies mosaic, mixup and rectangular training transformations to an item in the dataset.

    Args:
        index (int): Index of the item in the dataset.

    Returns:
        (dict): A dictionary containing the transformed item data.
    """
    labels = deepcopy(self.dataset[index])
    for transform in self.dataset.transforms.tolist():
        # Mosaic and mixup
        if hasattr(transform, 'get_indexes'):
            indexes = transform.get_indexes(self.dataset)
            if not isinstance(indexes, collections.abc.Sequence):
                indexes = [indexes]
            labels['mix_labels'] = [deepcopy(self.dataset[index]) for index in indexes]
        if self.dataset.rect and isinstance(transform, LetterBox):
            transform.new_shape = self.dataset.batch_shapes[self.dataset.batch[index]]
        labels = transform(labels)
        if 'mix_labels' in labels:
            labels.pop('mix_labels')
    return labels

__init__(dataset)

Parameters:

Name Type Description Default
dataset BaseDataset

The base dataset to apply transformations to.

required
Source code in ultralytics/yolo/data/dataset_wrappers.py
def __init__(self, dataset):
    """
    Args:
        dataset (BaseDataset): The base dataset to apply transformations to.
    """
    self.dataset = dataset
    self.imgsz = dataset.imgsz

__len__()

Returns the number of items in the dataset.

Source code in ultralytics/yolo/data/dataset_wrappers.py
def __len__(self):
    """Returns the number of items in the dataset."""
    return len(self.dataset)




Created 2023-04-16, Updated 2023-05-17
Authors: Glenn Jocher (3)