Menyiapkan Dataset COCO Kustom untuk Deteksi Objek dengan Python

Jan 21 2023

Deteksi objek adalah tugas penting dalam visi komputer dan memiliki berbagai aplikasi di berbagai bidang seperti mobil tanpa pengemudi, pengawasan video, dan pengambilan gambar. Salah satu dataset yang paling populer dan banyak digunakan untuk deteksi objek adalah dataset Common Objects in Context (COCO). Namun, terkadang kumpulan data COCO mungkin tidak berisi objek tertentu yang diminati atau mungkin memiliki terlalu banyak kelas untuk aplikasi tertentu. Dalam kasus seperti itu, membuat kumpulan data COCO khusus hanya dengan kelas dan gambar yang diperlukan dapat berguna.

Pada artikel ini, kita akan melalui proses pembuatan dataset COCO khusus untuk deteksi objek menggunakan Python. Kami akan menggunakan dataset COCO dan pustaka pycocotools untuk mengekstrak anotasi untuk kelas tertentu dan membuat file JSON baru yang berisi anotasi dan informasi gambar. Kami juga akan menunjukkan cara membuat tautan simbolis ke gambar asli di direktori baru.

Langkah 1: Instal Perpustakaan yang Diperlukan

Langkah pertama adalah menginstal pustaka yang diperlukan. Kami akan menggunakan perpustakaan berikut:

numpy: untuk bekerja dengan array dan matriks
json: untuk bekerja dengan file JSON
os: untuk bekerja dengan sistem file
pycocotools: untuk bekerja dengan kumpulan data COCO
tqdmuntuk membuat bilah kemajuan untuk langkah-langkah pemrosesan gambar

pip install numpy pycocotools tqdm

Langkah 2: Muat Kumpulan Data COCO dan Ekstrak Anotasi

Langkah selanjutnya adalah memuat set data COCO dan mengekstrak anotasi untuk kelas tertentu. Kami akan menggunakan COCOkelas dari pycocotoolsperpustakaan untuk memuat file JSON anotasi COCO. Kemudian, kita akan menggunakan getCatIdsmetode untuk mendapatkan ID kategori untuk kelas yang diminati dan getImgIdsmetode untuk mendapatkan ID gambar untuk gambar yang berisi kelas tersebut.

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

Sekarang setelah kita memiliki ID gambar untuk kelas "mobil" di set kereta dan validasi, kita dapat membuat subset dari set data COCO 2017 untuk kelas "mobil". Salah satu cara melakukannya adalah dengan membuat kumpulan direktori baru untuk menyimpan gambar dan anotasi untuk kelas "mobil". Kita bisa melakukan ini menggunakan perpustakaan 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)

Setelah kami memiliki ID gambar dan anotasi, kami dapat membuat file JSON baru yang berisi anotasi dan informasi gambar, dan menggunakan COCO API untuk memuat gambar dan menyimpannya ke direktori baru. Kami akan menggunakan loadImgsdan loadAnnsmetode dari COCOkelas untuk mendapatkan data gambar dan anotasi masing-masing, dan kemudian menggunakan pustaka json untuk menulis data ini ke file baru. Kami juga akan membuat tautan simbolis ke gambar asli di direktori baru, sehingga kumpulan data baru dapat digunakan dengan mudah dengan kerangka deteksi objek apa pun.

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)

Penting untuk diperhatikan bahwa, dalam contoh ini, skrip hanya mengekstraksi anotasi untuk kelas "mobil" dari kumpulan data COCO. Namun, skrip ini dapat dengan mudah dimodifikasi untuk mengekstraksi anotasi untuk kelas lain atau untuk mengekstraksi anotasi dari kumpulan data lain.

Kesimpulannya, membuat kumpulan data khusus dapat menjadi tugas yang memakan waktu, tetapi dengan alat dan teknik yang tepat, hal itu dapat dilakukan dengan relatif mudah. Dengan menggunakan API COCO dan pustaka json, kami dapat mengekstrak anotasi dan informasi gambar untuk kelas tertentu, dan membuat kumpulan data baru yang dapat digunakan untuk deteksi objek. Dengan kumpulan data baru ini, kami dapat melatih model deteksi objek khusus yang dapat disesuaikan dengan kasus penggunaan tertentu.