coco2yolo

Custom YoloV5 database creation with custom classes support

Based on PyCOCOtools API and COCO database

This project was meant to be used for YoloV5 database creation, but you can easily adapt it for any other model format.

Installation

pip install -r requirements.txt

Usage

You can use example.py and fit class arguments to your project, or import the class in your python project from coco2yolo import coco2yolo .

If you don't have any COCO Annotation file (instance_train2017.json), the program will download it automagically and use instance_train2017.json

Example

Call coco2yolo class with desired arguments

In this example, our data model has a custom class and the following classes indexes apply in labels:

0: Face mask, 1: car, 2: person, 3: bicycle

from coco2yolo import coco2yolo

coco2yolo(
    CUSTOM_DATASET=True,
    COCO_CATEGORIES_QUERY=['face mask', 'car', 'person', 'bicycle'],
    COCO_ANNOTATIONS="instances_train2017.json",
    NUM_SAMPLES=1500,
    OUTPUT_DIRECTORY="./output",
    CONVERT_YOLO=True
    )

Parameters

(list) COCO_CATEGORIES_QUERY : COCO categories to download (put abstract category to artificially increment cat_id index in your model)
(bool) CUSTOM_DATASET        : toggle custom classes indexes. Label indexes will be based on COCO_CATEGORIES_QUERY indexes
(str)  COCO_ANNOTATIONS      : path to COCO annotation json
(int)  NUM_SAMPLES           : number of images to download per category
(str)  OUTPUT_DIRECTORY      : path to output directory
(bool) CONVERT_YOLO          : convert bboxes to YOLOv5 format (from pixels location to image ratio location)

Important notes :
  - no label will be generated if you dont export to YoloV5 Format
  - if you don't have any COCO Annotation file (instance_train2017.json)
    -> the program will download it automagically and use instance_train2017.json
    -> the downloaded file location and name will be set by COCO_ANNOTATIONS argument

Output

../path/to/output/

        car/
            images/
                29327.jpg/
                ..
                8721.jpg/
            labels/
                29327.txt/
                ..
                8721.txt/
        
        person/
            images/
                761.jpg/
                ..
                91272.jpg/
            labels/
                761.txt/
                ..
                91272.txt/