Préparation d'un ensemble de données COCO personnalisé pour la détection d'objets avec Python

Jan 21 2023

La détection d'objets est une tâche cruciale en vision par ordinateur et a un large éventail d'applications dans des domaines tels que les voitures autonomes, la vidéosurveillance et la récupération d'images. L'un des ensembles de données les plus populaires et les plus utilisés pour la détection d'objets est l'ensemble de données Common Objects in Context (COCO). Cependant, il arrive parfois que l'ensemble de données COCO ne contienne pas les objets d'intérêt spécifiques ou ait trop de classes pour une application particulière. Dans de tels cas, la création d'un jeu de données COCO personnalisé avec uniquement les classes et les images nécessaires peut être utile.

Dans cet article, nous allons passer en revue le processus de création d'un ensemble de données COCO personnalisé pour la détection d'objets à l'aide de Python. Nous utiliserons l'ensemble de données COCO et la bibliothèque pycocotools pour extraire les annotations d'une classe spécifique et créer un nouveau fichier JSON contenant les annotations et les informations d'image. Nous montrerons également comment créer des liens symboliques vers les images originales dans un nouveau répertoire.

Étape 1 : Installez les bibliothèques requises

La première étape consiste à installer les bibliothèques requises. Nous utiliserons les bibliothèques suivantes :

numpy: pour travailler avec des tableaux et des matrices
json: pour travailler avec des fichiers JSON
os: pour travailler avec le système de fichiers
pycocotools: pour travailler avec le jeu de données COCO
tqdmpour créer une barre de progression pour les étapes de traitement d'image

pip install numpy pycocotools tqdm

Étape 2 : Charger l'ensemble de données COCO et extraire les annotations

L'étape suivante consiste à charger le jeu de données COCO et à extraire les annotations pour une classe spécifique. Nous allons utiliser la COCOclasse de la pycocotoolsbibliothèque pour charger le fichier JSON d'annotations COCO. Ensuite, nous utiliserons la getCatIdsméthode pour obtenir les ID de catégorie pour la classe d'intérêt et la getImgIdsméthode pour obtenir les ID d'image pour les images contenant la classe.

from pycocotools.coco import COCO

train_annotations = COCO("/path/to/instances_train2017.json")
valid_annotations = COCO("/path/to/instances_val2017.json")

# Get category IDs for "car"
cat_ids = train_annotations.getCatIds(catNms=["car"])

# Load category data for "car"
cats = train_annotations.loadCats(cat_ids)

# Get image IDs for "car" in train and val sets
train_img_ids = train_annotations.getImgIds(catIds=cat_ids)
valid_img_ids = valid_annotations.getImgIds(catIds=cat_ids)

# Print number of images
print(f"Number of training images: {len(train_img_ids)}")
print(f"Number of validation images: {len(valid_img_ids)}")

Maintenant que nous avons les ID d'image pour la classe "voiture" dans les ensembles de train et de validation, nous pouvons créer un sous-ensemble de l'ensemble de données COCO 2017 pour la classe "voiture". Une façon de faire est de créer un nouvel ensemble de répertoires pour stocker les images et les annotations pour la classe "voiture". Nous pouvons le faire en utilisant la bibliothèque os.

import os

# Define the path for the COCO directory
coco_subset_path = "/path/to/coco_subset"

# Create a list of directories to be created
dirs = ["train", "val", "annotations"]

# Loop through the list of directories and create them
for dir_name in dirs:
    dir_path = os.path.join(coco_subset_path, dir_name)
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

Une fois que nous avons les ID d'image et les annotations, nous pouvons créer un nouveau fichier JSON contenant les annotations et les informations d'image, et utiliser l'API COCO pour charger les images et les enregistrer dans les nouveaux répertoires. Nous utiliserons les méthodes loadImgset loadAnnsde la COCOclasse pour obtenir respectivement les données d'image et les annotations, puis utiliserons la bibliothèque json pour écrire ces données dans un nouveau fichier. Nous créerons également des liens symboliques vers les images d'origine dans un nouveau répertoire, afin que le nouvel ensemble de données puisse être utilisé facilement avec n'importe quel framework de détection d'objets.

import numpy as np
import json
from tqdm import tqdm
import logging

logging.basicConfig(level=logging.DEBUG)

def prepare_ds(root_path, _type, dest_path, img_ids, coco_annotations):
    # validate input parameters
    if not os.path.exists(root_path):
        raise ValueError(f"{root_path} does not exist")
    if _type not in ["train", "val", "test"]:
        raise ValueError(f"{_type} is not a valid value. Use 'train', 'test' or 'val'")
    if not os.path.exists(dest_path):
        os.makedirs(dest_path)
    if not os.path.exists(f'{dest_path}/annotations'):
        os.makedirs(f'{dest_path}/annotations')
        
    anns = []   
    imgs = []
    
    for img_ids in tqdm(img_ids):
        img_data = coco_annotations.loadImgs(img_ids)[0]
        imgs.append(img_data)
        
        img_file = str(root_path + '/' + img_data["file_name"])

        ann_ids = coco_annotations.getAnnIds(imgIds=img_data['id'], catIds=cat_ids,iscrowd=None)
        
        loaded_anns = coco_annotations.loadAnns(ann_ids)                
        for a in loaded_anns:
            anns.append(a)
        
        # using openCV for image processing 
        #mask = cv2.bitwise_or.reduce([train_annotations.annToMask(ann) * ann["category_id"] for ann in anns])
        
        # create symbolic link to original file in destination folder
        src = img_file        
        dest = f'{dest_path}/{_type}/{img_data["file_name"]}'        
        os.symlink(src, dest)
    
    data = {
         "info": coco_annotations.dataset['info'],
         "licenses": coco_annotations.dataset['licenses'],
         "images": imgs, 
         "annotations": anns,
         "categories": cats 
         }  
        
    dump = json.dumps(data)
        
    # save data to JSON file
    json_file = f'{dest_path}/annotations/instances_{_type}.json'
    try:
        with open(json_file, 'w') as f:
            f.write(str(dump))
            logging.info(f"Data saved to {json_file}")
    except Exception as e:
        logging.error(f"Error saving data to {json_file}: {e}")

def prepare_ds_and_anns(json_file_name, root_path, _type, dest_path, img_ids, coco_annotations):
    json_file = os.path.join(dest_path, "annotations", json_file_name)
    if os.path.exists(json_file):
        os.remove(json_file)
    print(f"{json_file} removed.")
    prepare_ds(root_path, _type, dest_path, img_ids, coco_annotations)

prepare_ds_and_anns("instances_train2017.json", '/path/to/coco2017/train2017', 'train', coco_subset_path, train_img_ids, train_annotations)

prepare_ds_and_anns("instances_val2017.json", '/path/to/coco2017/val2017', 'val', coco_subset_path, valid_img_ids, valid_annotations)

Il est important de noter que, dans cet exemple, le script extrait uniquement les annotations pour la classe « car » du jeu de données COCO. Cependant, ce script peut être facilement modifié pour extraire des annotations pour d'autres classes ou pour extraire des annotations d'autres jeux de données.

En conclusion, la création d'un ensemble de données personnalisé peut être une tâche fastidieuse, mais avec les bons outils et techniques, cela peut être fait relativement facilement. En utilisant l'API COCO et la bibliothèque json, nous avons pu extraire des annotations et des informations d'image pour une classe spécifique, et créer un nouvel ensemble de données pouvant être utilisé pour la détection d'objets. Avec ce nouvel ensemble de données, nous pouvons former des modèles de détection d'objets personnalisés qui peuvent être affinés pour s'adapter à des cas d'utilisation spécifiques.