0
# COCO Integration
1

2
SAHI provides comprehensive COCO dataset integration including loading, manipulation, annotation processing, evaluation, and format conversion capabilities. The framework supports the complete COCO ecosystem with seamless integration for both images and videos.
3

4
## Capabilities
5

6
### COCO Dataset Class
7

8
Main class for loading and manipulating COCO format datasets with full annotation support.
9

10
```python { .api }
11
class Coco:
12
    def __init__(self, coco_path: Optional[str] = None):
13
        """
14
        Initialize COCO dataset from JSON file.
15
        
16
        Parameters:
17
        - coco_path (str, optional): Path to COCO format JSON file
18
        """
19
    
20
    @property
21
    def json(self) -> Dict: ...
22
    
23
    @property
24
    def images(self) -> List[CocoImage]: ...
25
    
26
    @property
27
    def annotations(self) -> List[CocoAnnotation]: ...
28
    
29
    @property
30
    def categories(self) -> List[CocoCategory]: ...
31
    
32
    def add_image(self, coco_image: CocoImage) -> int:
33
        """
34
        Add image to dataset.
35
        
36
        Parameters:
37
        - coco_image: CocoImage instance
38
        
39
        Returns:
40
        int: Assigned image ID
41
        """
42
    
43
    def add_annotation(self, coco_annotation: CocoAnnotation) -> int:
44
        """
45
        Add annotation to dataset.
46
        
47
        Parameters:
48
        - coco_annotation: CocoAnnotation instance
49
        
50
        Returns:
51
        int: Assigned annotation ID
52
        """
53
    
54
    def add_category(self, coco_category: CocoCategory) -> int:
55
        """
56
        Add category to dataset.
57
        
58
        Parameters:
59
        - coco_category: CocoCategory instance
60
        
61
        Returns:
62
        int: Assigned category ID
63
        """
64
    
65
    def merge(self, coco2: "Coco") -> "Coco":
66
        """
67
        Merge with another COCO dataset.
68
        
69
        Parameters:
70
        - coco2: Another Coco instance to merge
71
        
72
        Returns:
73
        Coco: New merged dataset
74
        """
75
    
76
    def export_as_yolo(
77
        self,
78
        output_dir: str,
79
        train_split_rate: float = 1.0,
80
        numpy_seed: int = 0,
81
        mp: bool = True,
82
    ):
83
        """
84
        Export dataset in YOLO format.
85
        
86
        Parameters:
87
        - output_dir (str): Output directory for YOLO files
88
        - train_split_rate (float): Fraction for training set (0-1)
89
        - numpy_seed (int): Random seed for reproducible splits
90
        - mp (bool): Use multiprocessing for faster conversion
91
        """
92
    
93
    def save(self, save_path: str):
94
        """Save COCO dataset to JSON file."""
95
    
96
    def get_image_list(self) -> List[CocoImage]:
97
        """Get list of all images in dataset."""
98
    
99
    def get_category_list(self) -> List[CocoCategory]:
100
        """Get list of all categories in dataset."""
101
```
102

103
### COCO Image Class
104

105
Represents individual images in COCO datasets with annotation management.
106

107
```python { .api }
108
class CocoImage:
109
    def __init__(
110
        self, 
111
        image_path: str, 
112
        image_id: Optional[int] = None
113
    ):
114
        """
115
        Initialize COCO image.
116
        
117
        Parameters:
118
        - image_path (str): Path to image file
119
        - image_id (int, optional): Unique image identifier
120
        """
121
    
122
    @property
123
    def json(self) -> Dict: ...
124
    
125
    @property
126
    def annotations(self) -> List[CocoAnnotation]: ...
127
    
128
    @property
129
    def image_path(self) -> str: ...
130
    
131
    @property
132
    def file_name(self) -> str: ...
133
    
134
    def add_annotation(self, annotation: CocoAnnotation):
135
        """
136
        Add annotation to this image.
137
        
138
        Parameters:
139
        - annotation: CocoAnnotation instance
140
        """
141
    
142
    def get_annotation_list(self) -> List[CocoAnnotation]:
143
        """Get all annotations for this image."""
144
```
145

146
### COCO Annotation Class
147

148
Represents object annotations with bounding boxes and segmentation masks.
149

150
```python { .api }
151
class CocoAnnotation:
152
    def __init__(
153
        self,
154
        bbox: Optional[List[int]] = None,
155
        category_id: Optional[int] = None,
156
        category_name: Optional[str] = None,
157
        iscrowd: int = 0,
158
        area: Optional[int] = None,
159
        segmentation: Optional[List] = None,
160
        image_id: Optional[int] = None,
161
        annotation_id: Optional[int] = None,
162
    ):
163
        """
164
        Initialize COCO annotation.
165
        
166
        Parameters:
167
        - bbox (List[int], optional): Bounding box [x, y, width, height]
168
        - category_id (int, optional): Category identifier
169
        - category_name (str, optional): Category name
170
        - iscrowd (int): Whether annotation represents crowd (0 or 1)
171
        - area (int, optional): Annotation area in pixels
172
        - segmentation (List, optional): Polygon segmentation
173
        - image_id (int, optional): Associated image ID
174
        - annotation_id (int, optional): Unique annotation ID
175
        """
176
    
177
    @property
178
    def json(self) -> Dict: ...
179
    
180
    @property
181
    def area(self) -> int: ...
182
    
183
    @property
184
    def bbox(self) -> List[int]: ...
185
    
186
    def update_bbox_stats(self):
187
        """Update bounding box statistics from segmentation."""
188
```
189

190
### COCO Prediction Class
191

192
Represents model predictions in COCO format with confidence scores.
193

194
```python { .api }
195
class CocoPrediction:
196
    def __init__(
197
        self,
198
        bbox: List[int],
199
        category_id: int,
200
        score: float,
201
        category_name: Optional[str] = None,
202
        segmentation: Optional[List] = None,
203
        image_id: Optional[int] = None,
204
    ):
205
        """
206
        Initialize COCO prediction.
207
        
208
        Parameters:
209
        - bbox (List[int]): Bounding box [x, y, width, height]
210
        - category_id (int): Predicted category ID
211
        - score (float): Confidence score (0-1)
212
        - category_name (str, optional): Category name
213
        - segmentation (List, optional): Predicted segmentation mask
214
        - image_id (int, optional): Associated image ID
215
        """
216
    
217
    @property
218
    def json(self) -> Dict: ...
219
    
220
    @property
221
    def score(self) -> float: ...
222
    
223
    @property
224
    def category_id(self) -> int: ...
225
    
226
    @property
227
    def bbox(self) -> List[int]: ...
228
```
229

230
### COCO Category Class
231

232
Represents object categories in COCO datasets.
233

234
```python { .api }
235
class CocoCategory:
236
    def __init__(
237
        self, 
238
        category_id: int, 
239
        name: str, 
240
        supercategory: str = ""
241
    ):
242
        """
243
        Initialize COCO category.
244
        
245
        Parameters:
246
        - category_id (int): Unique category identifier
247
        - name (str): Category name
248
        - supercategory (str): Parent category name
249
        """
250
    
251
    @property
252
    def json(self) -> Dict: ...
253
    
254
    @property
255
    def id(self) -> int: ...
256
    
257
    @property
258
    def name(self) -> str: ...
259
```
260

261
### COCO Video Classes
262

263
Support for COCO-style video datasets with temporal annotations.
264

265
```python { .api }
266
class CocoVideo:
267
    def __init__(
268
        self, 
269
        video_path: str, 
270
        video_id: Optional[int] = None
271
    ):
272
        """
273
        Initialize COCO video.
274
        
275
        Parameters:
276
        - video_path (str): Path to video file
277
        - video_id (int, optional): Unique video identifier
278
        """
279

280
class CocoVidImage(CocoImage):
281
    """COCO video frame with temporal information."""
282

283
class CocoVidAnnotation(CocoAnnotation):
284
    """COCO video annotation with track information."""
285

286
class CocoVid(Coco):
287
    """COCO video dataset with temporal support."""
288
```
289

290
### Utility Functions
291

292
```python { .api }
293
def create_coco_dict() -> Dict:
294
    """
295
    Create empty COCO format dictionary structure.
296
    
297
    Returns:
298
    Dict: Empty COCO dictionary with required fields
299
    """
300

301
def merge(coco1: Coco, coco2: Coco) -> Coco:
302
    """
303
    Merge two COCO datasets.
304
    
305
    Parameters:
306
    - coco1: First COCO dataset
307
    - coco2: Second COCO dataset
308
    
309
    Returns:
310
    Coco: Merged dataset
311
    """
312

313
def export_coco_as_yolo(
314
    coco_path: str,
315
    output_dir: str,
316
    train_split_rate: float = 1.0,
317
    numpy_seed: int = 0,
318
) -> str:
319
    """
320
    Convert COCO dataset to YOLO format.
321
    
322
    Parameters:
323
    - coco_path (str): Path to COCO JSON file
324
    - output_dir (str): Output directory for YOLO files
325
    - train_split_rate (float): Training set split ratio
326
    - numpy_seed (int): Random seed for reproducible splits
327
    
328
    Returns:
329
    str: Path to output directory
330
    """
331
```
332

333
### Dataset Statistics
334

335
```python { .api }
336
class DatasetClassCounts:
337
    def __init__(self, coco: Coco):
338
        """
339
        Calculate dataset statistics and class distributions.
340
        
341
        Parameters:
342
        - coco: COCO dataset instance
343
        """
344
    
345
    @property
346
    def stats(self) -> Dict: ...
347
    
348
    def get_class_distribution(self) -> Dict[str, int]:
349
        """Get distribution of annotations per class."""
350
    
351
    def get_image_distribution(self) -> Dict[str, int]:
352
        """Get distribution of images per class."""
353
```
354

355
## Usage Examples
356

357
### Loading and Manipulating COCO Datasets
358

359
```python
360
from sahi.utils.coco import Coco, CocoImage, CocoAnnotation, CocoCategory
361

362
# Load existing COCO dataset
363
coco = Coco("dataset/annotations.json")
364

365
print(f"Dataset contains:")
366
print(f"  Images: {len(coco.images)}")
367
print(f"  Annotations: {len(coco.annotations)}")
368
print(f"  Categories: {len(coco.categories)}")
369

370
# Access images and annotations
371
for image in coco.images[:5]:  # First 5 images
372
    print(f"Image: {image.file_name}")
373
    print(f"  Annotations: {len(image.annotations)}")
374
    
375
    for annotation in image.annotations:
376
        print(f"    Category: {annotation.category_name}")
377
        print(f"    BBox: {annotation.bbox}")
378
```
379

380
### Creating New COCO Dataset
381

382
```python
383
from sahi.utils.coco import Coco, CocoImage, CocoAnnotation, CocoCategory
384

385
# Create new empty dataset
386
coco = Coco()
387

388
# Add categories
389
person_cat = CocoCategory(category_id=1, name="person")
390
car_cat = CocoCategory(category_id=2, name="car")
391
coco.add_category(person_cat)
392
coco.add_category(car_cat)
393

394
# Add image
395
image = CocoImage(image_path="images/sample.jpg", image_id=1)
396

397
# Add annotations
398
annotation1 = CocoAnnotation(
399
    bbox=[100, 100, 50, 80],  # [x, y, width, height]
400
    category_id=1,
401
    category_name="person",
402
    image_id=1,
403
    annotation_id=1
404
)
405

406
annotation2 = CocoAnnotation(
407
    bbox=[200, 150, 120, 60],
408
    category_id=2,
409
    category_name="car", 
410
    image_id=1,
411
    annotation_id=2
412
)
413

414
image.add_annotation(annotation1)
415
image.add_annotation(annotation2)
416
coco.add_image(image)
417

418
# Save dataset
419
coco.save("new_dataset.json")
420
```
421

422
### Converting to YOLO Format
423

424
```python
425
from sahi.utils.coco import export_coco_as_yolo
426

427
# Convert COCO to YOLO format
428
output_path = export_coco_as_yolo(
429
    coco_path="dataset/annotations.json",
430
    output_dir="yolo_dataset/",
431
    train_split_rate=0.8,  # 80% for training, 20% for validation
432
    numpy_seed=42          # For reproducible splits
433
)
434

435
print(f"YOLO dataset created at: {output_path}")
436
```
437

438
### Merging Datasets
439

440
```python
441
from sahi.utils.coco import Coco
442

443
# Load multiple datasets
444
dataset1 = Coco("dataset1/annotations.json")
445
dataset2 = Coco("dataset2/annotations.json")
446
dataset3 = Coco("dataset3/annotations.json")
447

448
# Merge datasets
449
merged = dataset1.merge(dataset2).merge(dataset3)
450

451
print(f"Merged dataset contains:")
452
print(f"  Images: {len(merged.images)}")
453
print(f"  Annotations: {len(merged.annotations)}")
454

455
# Save merged dataset
456
merged.save("merged_dataset.json")
457
```
458

459
### Working with Predictions
460

461
```python
462
from sahi.utils.coco import CocoPrediction
463
from sahi import get_sliced_prediction
464

465
# Get predictions from SAHI
466
result = get_sliced_prediction(
467
    image="test_image.jpg",
468
    detection_model=model
469
)
470

471
# Convert to COCO predictions
472
coco_predictions = result.to_coco_predictions()
473

474
# Access prediction details
475
for pred in coco_predictions:
476
    print(f"Category: {pred.category_name}")
477
    print(f"Confidence: {pred.score}")
478
    print(f"BBox: {pred.bbox}")
479

480
# Save predictions to JSON
481
import json
482
predictions_dict = [pred.json for pred in coco_predictions]
483
with open("predictions.json", "w") as f:
484
    json.dump(predictions_dict, f)
485
```
486

487
### Dataset Analysis
488

489
```python
490
from sahi.utils.coco import Coco, DatasetClassCounts
491

492
# Load dataset
493
coco = Coco("dataset/annotations.json")
494

495
# Analyze class distribution
496
stats = DatasetClassCounts(coco)
497
class_dist = stats.get_class_distribution()
498

499
print("Class distribution:")
500
for class_name, count in class_dist.items():
501
    print(f"  {class_name}: {count} annotations")
502

503
# Get images per class
504
image_dist = stats.get_image_distribution()
505
print("\nImages per class:")
506
for class_name, count in image_dist.items():
507
    print(f"  {class_name}: {count} images")
508
```
509

510
### Advanced Dataset Manipulation
511

512
```python
513
from sahi.utils.coco import Coco
514

515
# Load dataset
516
coco = Coco("large_dataset.json")
517

518
# Filter dataset by category
519
person_images = []
520
for image in coco.images:
521
    has_person = any(
522
        ann.category_name == "person" 
523
        for ann in image.annotations
524
    )
525
    if has_person:
526
        person_images.append(image)
527

528
print(f"Found {len(person_images)} images with people")
529

530
# Create subset dataset
531
subset_coco = Coco()
532
for image in person_images[:100]:  # First 100 images with people
533
    subset_coco.add_image(image)
534
    for annotation in image.annotations:
535
        subset_coco.add_annotation(annotation)
536

537
# Add categories
538
for category in coco.categories:
539
    subset_coco.add_category(category)
540

541
subset_coco.save("person_subset.json")
542
```