0
# Dataset Management
1

2
Tools for loading, processing, and converting datasets between popular computer vision formats like COCO, YOLO, and Pascal VOC. Provides unified interfaces for working with detection and classification datasets.
3

4
## Capabilities
5

6
### Base Dataset Class
7

8
Abstract base class that defines the common interface for all dataset types.
9

10
```python { .api }
11
class BaseDataset(ABC):
12
    """
13
    Abstract base class for all dataset types.
14
    """
15
    
16
    @abstractmethod
17
    def __len__(self) -> int:
18
        """Return the number of samples in the dataset."""
19
        
20
    @abstractmethod
21
    def split(
22
        self,
23
        split_ratio: float = 0.8,
24
        random_state: int | None = None,
25
        shuffle: bool = True
26
    ) -> tuple["BaseDataset", "BaseDataset"]:
27
        """Split the dataset into train and validation sets."""
28
```
29

30
### Detection Dataset
31

32
Comprehensive dataset class for object detection tasks with support for multiple annotation formats.
33

34
```python { .api }
35
class DetectionDataset(BaseDataset):
36
    """
37
    Contains information about a detection dataset. Handles lazy image loading
38
    and annotation retrieval, dataset splitting, conversions into multiple formats.
39

40
    Attributes:
41
        classes (list[str]): List containing dataset class names
42
        images (list[str] | dict[str, np.ndarray]): List of image paths or dictionary of loaded images
43
        annotations (dict[str, Detections]): Dictionary mapping image path to annotations
44
    """
45
    
46
    def __init__(
47
        self,
48
        classes: list[str],
49
        images: list[str] | dict[str, np.ndarray],
50
        annotations: dict[str, Detections]
51
    ) -> None: ...
52

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

56
    def __getitem__(self, i: int) -> tuple[str, np.ndarray, Detections]:
57
        """Get image path, image data, and annotations at index i."""
58

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

62
    def split(
63
        self,
64
        split_ratio: float = 0.8,
65
        random_state: int | None = None,
66
        shuffle: bool = True
67
    ) -> tuple["DetectionDataset", "DetectionDataset"]:
68
        """Split dataset into train and validation sets."""
69

70
    @classmethod
71
    def merge(cls, dataset_list: list["DetectionDataset"]) -> "DetectionDataset":
72
        """Merge multiple DetectionDataset instances into one."""
73

74
    @classmethod
75
    def from_pascal_voc(
76
        cls,
77
        images_directory_path: str,
78
        annotations_directory_path: str,
79
        force_masks: bool = False
80
    ) -> "DetectionDataset":
81
        """Create dataset from Pascal VOC format annotations."""
82

83
    def as_pascal_voc(
84
        self,
85
        images_directory_path: str | None = None,
86
        annotations_directory_path: str | None = None,
87
        min_image_area_percentage: float = 0.0,
88
        max_image_area_percentage: float = 1.0,
89
        approximation_percentage: float = 0.0
90
    ) -> None:
91
        """Export dataset to Pascal VOC format."""
92

93
    @classmethod
94
    def from_yolo(
95
        cls,
96
        images_directory_path: str,
97
        annotations_directory_path: str,
98
        data_yaml_path: str,
99
        force_masks: bool = False,
100
        is_obb: bool = False
101
    ) -> "DetectionDataset":
102
        """Create dataset from YOLO format annotations."""
103

104
    def as_yolo(
105
        self,
106
        images_directory_path: str | None = None,
107
        annotations_directory_path: str | None = None,
108
        data_yaml_path: str | None = None,
109
        min_image_area_percentage: float = 0.0,
110
        max_image_area_percentage: float = 1.0,
111
        approximation_percentage: float = 0.0
112
    ) -> None:
113
        """Export dataset to YOLO format."""
114

115
    @classmethod
116
    def from_coco(
117
        cls,
118
        images_directory_path: str,
119
        annotations_path: str,
120
        force_masks: bool = False
121
    ) -> "DetectionDataset":
122
        """Create dataset from COCO format annotations."""
123

124
    def as_coco(
125
        self,
126
        images_directory_path: str | None = None,
127
        annotations_path: str | None = None,
128
        min_image_area_percentage: float = 0.0,
129
        max_image_area_percentage: float = 1.0,
130
        approximation_percentage: float = 0.0
131
    ) -> None:
132
        """Export dataset to COCO format."""
133
```
134

135
### Classification Dataset
136

137
Dataset class for image classification tasks with folder structure support.
138

139
```python { .api }
140
class ClassificationDataset(BaseDataset):
141
    """
142
    Contains information about a classification dataset, handles lazy image
143
    loading, dataset splitting.
144

145
    Attributes:
146
        classes (list[str]): List containing dataset class names
147
        images (list[str] | dict[str, np.ndarray]): List of image paths or dictionary of images
148
        annotations (dict[str, Classifications]): Dictionary mapping image name to annotations
149
    """
150
    
151
    def __init__(
152
        self,
153
        classes: list[str],
154
        images: list[str] | dict[str, np.ndarray],
155
        annotations: dict[str, Classifications]
156
    ) -> None: ...
157

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

161
    def __getitem__(self, i: int) -> tuple[str, np.ndarray, Classifications]:
162
        """Get image path, image data, and classification at index i."""
163

164
    def __iter__(self) -> Iterator[tuple[str, np.ndarray, Classifications]]:
165
        """Iterate over images and classifications in the dataset."""
166

167
    def split(
168
        self,
169
        split_ratio: float = 0.8,
170
        random_state: int | None = None,
171
        shuffle: bool = True
172
    ) -> tuple["ClassificationDataset", "ClassificationDataset"]:
173
        """Split dataset into train and validation sets."""
174

175
    def as_folder_structure(self, root_directory_path: str) -> None:
176
        """Export dataset as folder structure with class subdirectories."""
177

178
    @classmethod
179
    def from_folder_structure(cls, root_directory_path: str) -> "ClassificationDataset":
180
        """Create dataset from folder structure with class subdirectories."""
181
```
182

183
### Dataset Utilities
184

185
Helper functions for dataset operations and format conversions.
186

187
```python { .api }
188
def get_coco_class_index_mapping(annotations_path: str) -> dict[int, str]:
189
    """Get mapping from COCO category IDs to class names."""
190

191
def mask_to_rle(mask: np.ndarray) -> dict[str, Any]:
192
    """Convert binary mask to COCO RLE format."""
193

194
def rle_to_mask(rle: dict[str, Any]) -> np.ndarray:
195
    """Convert COCO RLE format to binary mask."""
196
```
197

198
## Usage Examples
199

200
### Loading Different Dataset Formats
201

202
```python
203
import supervision as sv
204

205
# Load YOLO dataset
206
dataset = sv.DetectionDataset.from_yolo(
207
    images_directory_path="./images",
208
    annotations_directory_path="./labels", 
209
    data_yaml_path="./data.yaml"
210
)
211

212
# Load COCO dataset
213
coco_dataset = sv.DetectionDataset.from_coco(
214
    images_directory_path="./coco/images",
215
    annotations_path="./coco/annotations.json"
216
)
217

218
# Load Pascal VOC dataset
219
voc_dataset = sv.DetectionDataset.from_pascal_voc(
220
    images_directory_path="./images",
221
    annotations_directory_path="./annotations"
222
)
223

224
print(f"Dataset classes: {dataset.classes}")
225
print(f"Number of images: {len(dataset)}")
226
```
227

228
### Dataset Splitting and Processing
229

230
```python
231
import supervision as sv
232

233
# Load dataset
234
dataset = sv.DetectionDataset.from_yolo(
235
    images_directory_path="./images",
236
    annotations_directory_path="./labels",
237
    data_yaml_path="./data.yaml"
238
)
239

240
# Split into train and validation sets
241
train_dataset, val_dataset = dataset.split(
242
    split_ratio=0.8,
243
    random_state=42,
244
    shuffle=True
245
)
246

247
print(f"Train set: {len(train_dataset)} images")
248
print(f"Validation set: {len(val_dataset)} images")
249

250
# Iterate through dataset
251
for image_path, image, detections in train_dataset:
252
    print(f"Processing {image_path}")
253
    print(f"Found {len(detections)} objects")
254
    # Process image and annotations
255
```
256

257
### Format Conversions
258

259
```python
260
import supervision as sv
261

262
# Load YOLO dataset
263
dataset = sv.DetectionDataset.from_yolo(
264
    images_directory_path="./yolo/images",
265
    annotations_directory_path="./yolo/labels",
266
    data_yaml_path="./yolo/data.yaml"
267
)
268

269
# Convert to COCO format
270
dataset.as_coco(
271
    images_directory_path="./coco/images",
272
    annotations_path="./coco/annotations.json"
273
)
274

275
# Convert to Pascal VOC format
276
dataset.as_pascal_voc(
277
    images_directory_path="./voc/images",
278
    annotations_directory_path="./voc/annotations"
279
)
280

281
print("Dataset converted to multiple formats!")
282
```
283

284
### Merging Multiple Datasets
285

286
```python
287
import supervision as sv
288

289
# Load multiple datasets
290
dataset1 = sv.DetectionDataset.from_yolo(
291
    images_directory_path="./dataset1/images",
292
    annotations_directory_path="./dataset1/labels",
293
    data_yaml_path="./dataset1/data.yaml"
294
)
295

296
dataset2 = sv.DetectionDataset.from_coco(
297
    images_directory_path="./dataset2/images", 
298
    annotations_path="./dataset2/annotations.json"
299
)
300

301
# Merge datasets
302
merged_dataset = sv.DetectionDataset.merge([dataset1, dataset2])
303

304
print(f"Merged dataset: {len(merged_dataset)} images")
305
print(f"Combined classes: {merged_dataset.classes}")
306
```
307

308
### Classification Dataset Example
309

310
```python
311
import supervision as sv
312

313
# Load classification dataset from folder structure
314
# Expected structure:
315
# root/
316
#   ├── class1/
317
#   │   ├── img1.jpg
318
#   │   └── img2.jpg
319
#   └── class2/
320
#       ├── img3.jpg
321
#       └── img4.jpg
322

323
classification_dataset = sv.ClassificationDataset.from_folder_structure(
324
    root_directory_path="./classification_data"
325
)
326

327
# Split the dataset
328
train_cls, val_cls = classification_dataset.split(split_ratio=0.7)
329

330
# Export back to folder structure
331
train_cls.as_folder_structure("./output/train")
332
val_cls.as_folder_structure("./output/val")
333

334
# Iterate through classifications
335
for image_path, image, classification in classification_dataset:
336
    print(f"Image: {image_path}")
337
    print(f"Class ID: {classification.class_id}")
338
    print(f"Confidence: {classification.confidence}")
339
```