Skip to content

Core

Warning

Dataset API is still fluid and may change. If you use Dataset in your project until further notice, freeze the supervision version in your requirements.txt.

Dataset

Dataclass containing information about the dataset.

Attributes:

Name Type Description
classes List[str]

List containing dataset class names.
images Dict[str, ndarray]

Dictionary mapping image name to image.
annotations Dict[str, Detections]

Dictionary mapping image name to annotations.
Source code in supervision/dataset/core.py 
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
@dataclass
class Dataset:
    """
    Dataclass containing information about the dataset.

    Attributes:
        classes (List[str]): List containing dataset class names.
        images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
        annotations (Dict[str, Detections]): Dictionary mapping image name to annotations.
    """

    classes: List[str]
    images: Dict[str, np.ndarray]
    annotations: Dict[str, Detections]

    def __len__(self) -> int:
        """
        Return the number of images in the dataset.

        Returns:
            int: The number of images.
        """
        return len(self.images)

    def __iter__(self) -> Iterator[Tuple[str, np.ndarray, Detections]]:
        """
        Iterate over the images and annotations in the dataset.

        Yields:
            Iterator[Tuple[str, np.ndarray, Detections]]: An iterator that yields tuples containing the image name,
                                                          the image data, and its corresponding annotation.
        """
        for image_name, image in self.images.items():
            yield image_name, image, self.annotations.get(image_name, None)

    def as_pascal_voc(
        self,
        images_directory_path: Optional[str] = None,
        annotations_directory_path: Optional[str] = None,
        min_image_area_percentage: float = 0.0,
        max_image_area_percentage: float = 1.0,
        approximation_percentage: float = 0.75,
    ) -> None:
        """
        Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
        PASCAL VOC format, which consists of XML files. The method allows filtering the detections based on their area
        percentage.

        Args:
            images_directory_path (Optional[str]): The path to the directory where the images should be saved.
                If not provided, images will not be saved.
            annotations_directory_path (Optional[str]): The path to the directory where the annotations in
                PASCAL VOC format should be saved. If not provided, annotations will not be saved.
            min_image_area_percentage (float): The minimum percentage of detection area relative to
                the image area for a detection to be included.
            max_image_area_percentage (float): The maximum percentage of detection area relative to
                the image area for a detection to be included.
            approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
        """
        if images_directory_path:
            images_path = Path(images_directory_path)
            images_path.mkdir(parents=True, exist_ok=True)

        if annotations_directory_path:
            annotations_path = Path(annotations_directory_path)
            annotations_path.mkdir(parents=True, exist_ok=True)

        for image_name, image in self.images.items():
            detections = self.annotations[image_name]

            if images_directory_path:
                cv2.imwrite(str(images_path / image_name), image)

            if annotations_directory_path:
                annotation_name = Path(image_name).stem
                pascal_voc_xml = detections_to_pascal_voc(
                    detections=detections,
                    classes=self.classes,
                    filename=image_name,
                    image_shape=image.shape,
                    min_image_area_percentage=min_image_area_percentage,
                    max_image_area_percentage=max_image_area_percentage,
                    approximation_percentage=approximation_percentage,
                )

                with open(annotations_path / f"{annotation_name}.xml", "w") as f:
                    f.write(pascal_voc_xml)

    @classmethod
    def from_pascal_voc(
        cls, images_directory_path: str, annotations_directory_path: str
    ) -> Dataset:
        """
        Creates a Dataset instance from PASCAL VOC formatted data.

        Args:
            images_directory_path (str): The path to the directory containing the images.
            annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations.

        Returns:
            Dataset: A Dataset instance containing the loaded images and annotations.

        Example:
            ```python
            >>> import roboflow
            >>> from roboflow import Roboflow
            >>> import supervision as sv

            >>> roboflow.login()

            >>> rf = Roboflow()

            >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
            >>> dataset = project.version(PROJECT_VERSION).download("voc")

            >>> train_dataset = sv.Dataset.from_yolo(
            ...     images_directory_path=f"{dataset.location}/train/images",
            ...     annotations_directory_path=f"{dataset.location}/train/labels"
            ... )

            >>> dataset.classes
            ['dog', 'person']
            ```
        """
        image_paths = list_files_with_extensions(
            directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
        )
        annotation_paths = list_files_with_extensions(
            directory=annotations_directory_path, extensions=["xml"]
        )

        raw_annotations: List[Tuple[str, Detections, List[str]]] = [
            load_pascal_voc_annotations(annotation_path=str(annotation_path))
            for annotation_path in annotation_paths
        ]

        classes = []
        for annotation in raw_annotations:
            classes.extend(annotation[2])
        classes = list(set(classes))

        for annotation in raw_annotations:
            class_id = [classes.index(class_name) for class_name in annotation[2]]
            annotation[1].class_id = np.array(class_id)

        images = {
            image_path.name: cv2.imread(str(image_path)) for image_path in image_paths
        }

        annotations = {
            image_name: detections for image_name, detections, _ in raw_annotations
        }
        return Dataset(classes=classes, images=images, annotations=annotations)

    @classmethod
    def from_yolo(
        cls,
        images_directory_path: str,
        annotations_directory_path: str,
        data_yaml_path: str,
        force_masks: bool = False,
    ) -> Dataset:
        """
        Creates a Dataset instance from YOLO formatted data.

        Args:
            images_directory_path (str): The path to the directory containing the images.
            annotations_directory_path (str): The path to the directory containing the YOLO annotation files.
            data_yaml_path (str): The path to the data YAML file containing class information.
            force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.

        Returns:
            Dataset: A Dataset instance containing the loaded images and annotations.

        Example:
            ```python
            >>> import roboflow
            >>> from roboflow import Roboflow
            >>> import supervision as sv

            >>> roboflow.login()

            >>> rf = Roboflow()

            >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
            >>> dataset = project.version(PROJECT_VERSION).download("yolov5")

            >>> train_dataset = sv.Dataset.from_yolo(
            ...     images_directory_path=f"{dataset.location}/train/images",
            ...     annotations_directory_path=f"{dataset.location}/train/labels",
            ...     data_yaml_path=f"{dataset.location}/data.yaml"
            ... )

            >>> dataset.classes
            ['dog', 'person']
            ```
        """
        classes, images, annotations = load_yolo_annotations(
            images_directory_path=images_directory_path,
            annotations_directory_path=annotations_directory_path,
            data_yaml_path=data_yaml_path,
            force_masks=force_masks,
        )
        return Dataset(classes=classes, images=images, annotations=annotations)

__iter__()

Iterate over the images and annotations in the dataset.

Yields:

Type Description
str

Iterator[Tuple[str, np.ndarray, Detections]]: An iterator that yields tuples containing the image name, the image data, and its corresponding annotation.
Source code in supervision/dataset/core.py 
43
44
45
46
47
48
49
50
51
52
def __iter__(self) -> Iterator[Tuple[str, np.ndarray, Detections]]:
    """
    Iterate over the images and annotations in the dataset.

    Yields:
        Iterator[Tuple[str, np.ndarray, Detections]]: An iterator that yields tuples containing the image name,
                                                      the image data, and its corresponding annotation.
    """
    for image_name, image in self.images.items():
        yield image_name, image, self.annotations.get(image_name, None)

__len__()

Return the number of images in the dataset.

Returns:

Name Type Description
int int

The number of images.
Source code in supervision/dataset/core.py 
34
35
36
37
38
39
40
41
def __len__(self) -> int:
    """
    Return the number of images in the dataset.

    Returns:
        int: The number of images.
    """
    return len(self.images)

as_pascal_voc(images_directory_path=None, annotations_directory_path=None, min_image_area_percentage=0.0, max_image_area_percentage=1.0, approximation_percentage=0.75)

Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in PASCAL VOC format, which consists of XML files. The method allows filtering the detections based on their area percentage.

Parameters:

Name Type Description Default
images_directory_path Optional[str]

The path to the directory where the images should be saved. If not provided, images will not be saved.

 None
annotations_directory_path Optional[str]

The path to the directory where the annotations in PASCAL VOC format should be saved. If not provided, annotations will not be saved.

 None
min_image_area_percentage float

The minimum percentage of detection area relative to the image area for a detection to be included.

 0.0
max_image_area_percentage float

The maximum percentage of detection area relative to the image area for a detection to be included.

 1.0
approximation_percentage float

The percentage of polygon points to be removed from the input polygon, in the range [0, 1).

 0.75
Source code in supervision/dataset/core.py 
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def as_pascal_voc(
    self,
    images_directory_path: Optional[str] = None,
    annotations_directory_path: Optional[str] = None,
    min_image_area_percentage: float = 0.0,
    max_image_area_percentage: float = 1.0,
    approximation_percentage: float = 0.75,
) -> None:
    """
    Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
    PASCAL VOC format, which consists of XML files. The method allows filtering the detections based on their area
    percentage.

    Args:
        images_directory_path (Optional[str]): The path to the directory where the images should be saved.
            If not provided, images will not be saved.
        annotations_directory_path (Optional[str]): The path to the directory where the annotations in
            PASCAL VOC format should be saved. If not provided, annotations will not be saved.
        min_image_area_percentage (float): The minimum percentage of detection area relative to
            the image area for a detection to be included.
        max_image_area_percentage (float): The maximum percentage of detection area relative to
            the image area for a detection to be included.
        approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
    """
    if images_directory_path:
        images_path = Path(images_directory_path)
        images_path.mkdir(parents=True, exist_ok=True)

    if annotations_directory_path:
        annotations_path = Path(annotations_directory_path)
        annotations_path.mkdir(parents=True, exist_ok=True)

    for image_name, image in self.images.items():
        detections = self.annotations[image_name]

        if images_directory_path:
            cv2.imwrite(str(images_path / image_name), image)

        if annotations_directory_path:
            annotation_name = Path(image_name).stem
            pascal_voc_xml = detections_to_pascal_voc(
                detections=detections,
                classes=self.classes,
                filename=image_name,
                image_shape=image.shape,
                min_image_area_percentage=min_image_area_percentage,
                max_image_area_percentage=max_image_area_percentage,
                approximation_percentage=approximation_percentage,
            )

            with open(annotations_path / f"{annotation_name}.xml", "w") as f:
                f.write(pascal_voc_xml)

from_pascal_voc(images_directory_path, annotations_directory_path) classmethod

Creates a Dataset instance from PASCAL VOC formatted data.

Parameters:

Name Type Description Default
images_directory_path str

The path to the directory containing the images.

 required
annotations_directory_path str

The path to the directory containing the PASCAL VOC XML annotations.

 required

Returns:

Name Type Description
Dataset Dataset

A Dataset instance containing the loaded images and annotations.
Example 
>>> import roboflow
>>> from roboflow import Roboflow
>>> import supervision as sv

>>> roboflow.login()

>>> rf = Roboflow()

>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
>>> dataset = project.version(PROJECT_VERSION).download("voc")

>>> train_dataset = sv.Dataset.from_yolo(
...     images_directory_path=f"{dataset.location}/train/images",
...     annotations_directory_path=f"{dataset.location}/train/labels"
... )

>>> dataset.classes
['dog', 'person']
Source code in supervision/dataset/core.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
@classmethod
def from_pascal_voc(
    cls, images_directory_path: str, annotations_directory_path: str
) -> Dataset:
    """
    Creates a Dataset instance from PASCAL VOC formatted data.

    Args:
        images_directory_path (str): The path to the directory containing the images.
        annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations.

    Returns:
        Dataset: A Dataset instance containing the loaded images and annotations.

    Example:
        ```python
        >>> import roboflow
        >>> from roboflow import Roboflow
        >>> import supervision as sv

        >>> roboflow.login()

        >>> rf = Roboflow()

        >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
        >>> dataset = project.version(PROJECT_VERSION).download("voc")

        >>> train_dataset = sv.Dataset.from_yolo(
        ...     images_directory_path=f"{dataset.location}/train/images",
        ...     annotations_directory_path=f"{dataset.location}/train/labels"
        ... )

        >>> dataset.classes
        ['dog', 'person']
        ```
    """
    image_paths = list_files_with_extensions(
        directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
    )
    annotation_paths = list_files_with_extensions(
        directory=annotations_directory_path, extensions=["xml"]
    )

    raw_annotations: List[Tuple[str, Detections, List[str]]] = [
        load_pascal_voc_annotations(annotation_path=str(annotation_path))
        for annotation_path in annotation_paths
    ]

    classes = []
    for annotation in raw_annotations:
        classes.extend(annotation[2])
    classes = list(set(classes))

    for annotation in raw_annotations:
        class_id = [classes.index(class_name) for class_name in annotation[2]]
        annotation[1].class_id = np.array(class_id)

    images = {
        image_path.name: cv2.imread(str(image_path)) for image_path in image_paths
    }

    annotations = {
        image_name: detections for image_name, detections, _ in raw_annotations
    }
    return Dataset(classes=classes, images=images, annotations=annotations)

from_yolo(images_directory_path, annotations_directory_path, data_yaml_path, force_masks=False) classmethod

Creates a Dataset instance from YOLO formatted data.

Parameters:

Name Type Description Default
images_directory_path str

The path to the directory containing the images.

 required
annotations_directory_path str

The path to the directory containing the YOLO annotation files.

 required
data_yaml_path str

The path to the data YAML file containing class information.

 required
force_masks bool

If True, forces masks to be loaded for all annotations, regardless of whether they are present.

 False

Returns:

Name Type Description
Dataset Dataset

A Dataset instance containing the loaded images and annotations.
Example 
>>> import roboflow
>>> from roboflow import Roboflow
>>> import supervision as sv

>>> roboflow.login()

>>> rf = Roboflow()

>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
>>> dataset = project.version(PROJECT_VERSION).download("yolov5")

>>> train_dataset = sv.Dataset.from_yolo(
...     images_directory_path=f"{dataset.location}/train/images",
...     annotations_directory_path=f"{dataset.location}/train/labels",
...     data_yaml_path=f"{dataset.location}/data.yaml"
... )

>>> dataset.classes
['dog', 'person']
Source code in supervision/dataset/core.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
@classmethod
def from_yolo(
    cls,
    images_directory_path: str,
    annotations_directory_path: str,
    data_yaml_path: str,
    force_masks: bool = False,
) -> Dataset:
    """
    Creates a Dataset instance from YOLO formatted data.

    Args:
        images_directory_path (str): The path to the directory containing the images.
        annotations_directory_path (str): The path to the directory containing the YOLO annotation files.
        data_yaml_path (str): The path to the data YAML file containing class information.
        force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.

    Returns:
        Dataset: A Dataset instance containing the loaded images and annotations.

    Example:
        ```python
        >>> import roboflow
        >>> from roboflow import Roboflow
        >>> import supervision as sv

        >>> roboflow.login()

        >>> rf = Roboflow()

        >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
        >>> dataset = project.version(PROJECT_VERSION).download("yolov5")

        >>> train_dataset = sv.Dataset.from_yolo(
        ...     images_directory_path=f"{dataset.location}/train/images",
        ...     annotations_directory_path=f"{dataset.location}/train/labels",
        ...     data_yaml_path=f"{dataset.location}/data.yaml"
        ... )

        >>> dataset.classes
        ['dog', 'person']
        ```
    """
    classes, images, annotations = load_yolo_annotations(
        images_directory_path=images_directory_path,
        annotations_directory_path=annotations_directory_path,
        data_yaml_path=data_yaml_path,
        force_masks=force_masks,
    )
    return Dataset(classes=classes, images=images, annotations=annotations)