Preparando um conjunto de dados COCO personalizado para detecção de objetos com Python

Jan 21 2023

A detecção de objetos é uma tarefa crucial em visão computacional e possui uma ampla gama de aplicações em áreas como carros autônomos, vigilância por vídeo e recuperação de imagens. Um dos conjuntos de dados mais populares e amplamente utilizados para detecção de objetos é o conjunto de dados COCO (Common Objects in Context). No entanto, às vezes, o conjunto de dados COCO pode não conter os objetos de interesse específicos ou pode ter muitas classes para um aplicativo específico. Nesses casos, pode ser útil criar um conjunto de dados COCO personalizado apenas com as classes e imagens necessárias.

Neste artigo, veremos o processo de criação de um conjunto de dados COCO personalizado para detecção de objetos usando Python. Usaremos o conjunto de dados COCO e a biblioteca pycocotools para extrair anotações para uma classe específica e criar um novo arquivo JSON contendo as anotações e as informações da imagem. Também mostraremos como criar links simbólicos para as imagens originais em um novo diretório.

Etapa 1: instalar as bibliotecas necessárias

O primeiro passo é instalar as bibliotecas necessárias. Estaremos utilizando as seguintes bibliotecas:

numpy: para trabalhar com arrays e matrizes
json: para trabalhar com arquivos JSON
os: para trabalhar com o sistema de arquivos
pycocotools: para trabalhar com o conjunto de dados COCO
tqdmpara criar uma barra de progresso para etapas de processamento de imagem

pip install numpy pycocotools tqdm

Etapa 2: carregar o conjunto de dados COCO e extrair anotações

A próxima etapa é carregar o conjunto de dados COCO e extrair anotações para uma classe específica. Usaremos a COCOclasse da pycocotoolsbiblioteca para carregar o arquivo JSON de anotações do COCO. Em seguida, usaremos o getCatIdsmétodo para obter os IDs de categoria para a classe de interesse e o getImgIdsmétodo para obter os IDs de imagem para as imagens que contêm a 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)}")

Agora que temos os IDs de imagem para a classe “carro” nos conjuntos de trem e validação, podemos criar um subconjunto do conjunto de dados COCO 2017 para a classe “carro”. Uma maneira de fazer isso é criar um novo conjunto de diretórios para armazenar as imagens e anotações da classe “carro”. Podemos fazer isso usando a biblioteca 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)

Assim que tivermos os IDs e anotações da imagem, podemos criar um novo arquivo JSON contendo as anotações e informações da imagem e usar a API COCO para carregar as imagens e salvá-las nos novos diretórios. Usaremos os métodos loadImgse loadAnnsda COCOclasse para obter os dados da imagem e as anotações, respectivamente, e então usaremos a biblioteca json para gravar esses dados em um novo arquivo. Também criaremos links simbólicos para as imagens originais em um novo diretório, para que o novo conjunto de dados possa ser usado facilmente com qualquer estrutura de detecção de objetos.

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)

É importante observar que, neste exemplo, o script apenas extrai anotações para a classe “carro” do conjunto de dados COCO. No entanto, esse script pode ser facilmente modificado para extrair anotações para outras classes ou para extrair anotações de outros conjuntos de dados.

Concluindo, criar um conjunto de dados personalizado pode ser uma tarefa demorada, mas com as ferramentas e técnicas certas, isso pode ser feito com relativa facilidade. Usando a API COCO e a biblioteca json, conseguimos extrair anotações e informações de imagem para uma classe específica e criar um novo conjunto de dados que pode ser usado para detecção de objetos. Com esse novo conjunto de dados, podemos treinar modelos personalizados de detecção de objetos que podem ser ajustados para atender a casos de uso específicos.