Vorbereiten eines benutzerdefinierten COCO-Datensatzes für die Objekterkennung mit Python

Jan 21 2023

Die Objekterkennung ist eine entscheidende Aufgabe in der Computer Vision und hat ein breites Anwendungsspektrum in Bereichen wie selbstfahrenden Autos, Videoüberwachung und Bildabruf. Einer der beliebtesten und am weitesten verbreiteten Datensätze zur Objekterkennung ist der Common Objects in Context (COCO)-Datensatz. Manchmal enthält der COCO-Datensatz jedoch möglicherweise nicht die spezifischen Objekte von Interesse oder verfügt möglicherweise über zu viele Klassen für eine bestimmte Anwendung. In solchen Fällen kann es hilfreich sein, einen benutzerdefinierten COCO-Datensatz zu erstellen, der nur die erforderlichen Klassen und Bilder enthält.

In diesem Artikel werden wir den Prozess der Erstellung eines benutzerdefinierten COCO-Datensatzes für die Objekterkennung mit Python durchgehen. Wir werden den COCO-Datensatz und die Pycocotools-Bibliothek verwenden, um Anmerkungen für eine bestimmte Klasse zu extrahieren und eine neue JSON-Datei zu erstellen, die die Anmerkungen und Bildinformationen enthält. Außerdem zeigen wir, wie man symbolische Links zu den Originalbildern in einem neuen Verzeichnis erstellt.

Schritt 1: Installieren Sie die erforderlichen Bibliotheken

Der erste Schritt besteht darin, die erforderlichen Bibliotheken zu installieren. Wir werden die folgenden Bibliotheken verwenden:

numpy: zum Arbeiten mit Arrays und Matrizen
json: zum Arbeiten mit JSON-Dateien
os: für die Arbeit mit dem Dateisystem
pycocotools: für die Arbeit mit dem COCO-Datensatz
tqdmzum Erstellen eines Fortschrittsbalkens für Bildverarbeitungsschritte

pip install numpy pycocotools tqdm

Schritt 2: Laden Sie den COCO-Datensatz und extrahieren Sie Anmerkungen

Der nächste Schritt besteht darin, den COCO-Datensatz zu laden und Anmerkungen für eine bestimmte Klasse zu extrahieren. Wir werden die COCOKlasse aus der pycocotoolsBibliothek verwenden, um die COCO-Annotations-JSON-Datei zu laden. Anschließend verwenden wir die getCatIdsMethode, um die Kategorie-IDs für die interessierende Klasse abzurufen, und die getImgIdsMethode, um die Bild-IDs für Bilder abzurufen, die die Klasse enthalten.

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)}")

Nachdem wir nun die Bild-IDs für die Klasse „Auto“ in den Zug- und Validierungssätzen haben, können wir eine Teilmenge des COCO 2017-Datensatzes für die Klasse „Auto“ erstellen. Eine Möglichkeit, dies zu erreichen, besteht darin, einen neuen Satz von Verzeichnissen zu erstellen, um die Bilder und Anmerkungen für die Klasse „Auto“ zu speichern. Wir können dies mithilfe der OS-Bibliothek tun.

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)

Sobald wir die Bild-IDs und Anmerkungen haben, können wir eine neue JSON-Datei mit den Anmerkungen und Bildinformationen erstellen und die COCO-API verwenden, um die Bilder zu laden und in den neuen Verzeichnissen zu speichern. Wir verwenden die Methoden loadImgsund loadAnnsder COCOKlasse, um die Bilddaten bzw. Anmerkungen abzurufen, und verwenden dann die JSON-Bibliothek, um diese Daten in eine neue Datei zu schreiben. Wir werden außerdem symbolische Links zu den Originalbildern in einem neuen Verzeichnis erstellen, sodass der neue Datensatz problemlos mit jedem Objekterkennungs-Framework verwendet werden kann.

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)

Es ist wichtig zu beachten, dass das Skript in diesem Beispiel nur Anmerkungen für die Klasse „Auto“ aus dem COCO-Datensatz extrahiert. Dieses Skript kann jedoch leicht geändert werden, um Anmerkungen für andere Klassen oder Anmerkungen aus anderen Datensätzen zu extrahieren.

Zusammenfassend lässt sich sagen, dass die Erstellung eines benutzerdefinierten Datensatzes eine zeitaufwändige Aufgabe sein kann, mit den richtigen Tools und Techniken jedoch relativ einfach zu bewerkstelligen ist. Mithilfe der COCO-API und der JSON-Bibliothek konnten wir Anmerkungen und Bildinformationen für eine bestimmte Klasse extrahieren und einen neuen Datensatz erstellen, der zur Objekterkennung verwendet werden kann. Mit diesem neuen Datensatz können wir benutzerdefinierte Objekterkennungsmodelle trainieren, die an bestimmte Anwendungsfälle angepasst werden können.