Python으로 개체 감지를 위한 사용자 지정 COCO 데이터 세트 준비

Jan 21 2023

객체 감지는 컴퓨터 비전에서 중요한 작업이며 자율 주행 자동차, 비디오 감시 및 이미지 검색과 같은 영역에서 광범위한 응용 프로그램을 가지고 있습니다. 개체 감지를 위해 가장 인기 있고 널리 사용되는 데이터 세트 중 하나는 COCO(Common Objects in Context) 데이터 세트입니다.

객체 감지는 컴퓨터 비전에서 중요한 작업이며 자율 주행 자동차, 비디오 감시 및 이미지 검색과 같은 영역에서 광범위한 응용 프로그램을 가지고 있습니다. 개체 감지를 위해 가장 인기 있고 널리 사용되는 데이터 세트 중 하나는 COCO(Common Objects in Context) 데이터 세트입니다. 그러나 때로는 COCO 데이터 세트에 관심 있는 특정 개체가 포함되어 있지 않거나 특정 애플리케이션에 대해 너무 많은 클래스가 있을 수 있습니다. 이러한 경우 필요한 클래스와 이미지만 있는 사용자 지정 COCO 데이터 세트를 만드는 것이 유용할 수 있습니다.

이 기사에서는 Python을 사용하여 객체 감지를 위한 사용자 지정 COCO 데이터 세트를 만드는 과정을 살펴보겠습니다. COCO 데이터 세트와 pycocotools 라이브러리를 사용하여 특정 클래스에 대한 주석을 추출하고 주석과 이미지 정보가 포함된 새 JSON 파일을 생성합니다. 또한 새 디렉토리의 원본 이미지에 대한 심볼릭 링크를 만드는 방법도 보여줍니다.

1단계: 필수 라이브러리 설치

첫 번째 단계는 필요한 라이브러리를 설치하는 것입니다. 우리는 다음 라이브러리를 사용할 것입니다.

numpy: 배열 및 행렬 작업용
json: JSON 파일 작업용
os: 파일 시스템 작업용
pycocotools: COCO 데이터 세트 작업용
tqdm이미지 처리 단계에 대한 진행률 표시줄 생성

pip install numpy pycocotools tqdm

2단계: COCO 데이터 세트 로드 및 주석 추출

다음 단계는 COCO 데이터 세트를 로드하고 특정 클래스에 대한 주석을 추출하는 것입니다. COCO라이브러리 의 클래스를 사용하여 pycocotoolsCOCO 주석 JSON 파일을 로드합니다. getCatIds그런 다음 관심 있는 클래스의 카테고리 ID를 가져오는 메서드와 getImgIds클래스를 포함하는 이미지의 이미지 ID를 가져오는 메서드를 사용합니다 .

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

이제 기차 및 유효성 검사 세트에 "car" 클래스에 대한 이미지 ID가 있으므로 "car" 클래스에 대한 COCO 2017 데이터 세트의 하위 집합을 만들 수 있습니다. 이를 수행하는 한 가지 방법은 "car" 클래스의 이미지와 주석을 저장할 새 디렉토리 세트를 만드는 것입니다. 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)

이미지 ID와 주석이 있으면 주석과 이미지 정보가 포함된 새 JSON 파일을 만들고 COCO API를 사용하여 이미지를 로드하고 새 디렉터리에 저장할 수 있습니다. 클래스의 loadImgs및 메서드를 사용하여 각각 이미지 데이터와 주석을 가져온 다음 json 라이브러리를 사용하여 이 데이터를 새 파일에 씁니다. 또한 새 디렉터리에 원본 이미지에 대한 심볼릭 링크를 생성하여 새 데이터 세트를 모든 객체 감지 프레임워크와 함께 쉽게 사용할 수 있습니다.loadAnnsCOCO

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)

이 예에서 스크립트는 COCO 데이터 세트에서 "car" 클래스에 대한 주석만 추출한다는 점에 유의해야 합니다. 그러나 이 스크립트는 다른 클래스에 대한 주석을 추출하거나 다른 데이터 세트에서 주석을 추출하도록 쉽게 수정할 수 있습니다.

결론적으로 사용자 지정 데이터 세트를 만드는 것은 시간이 많이 걸리는 작업일 수 있지만 올바른 도구와 기술을 사용하면 비교적 쉽게 수행할 수 있습니다. COCO API와 json 라이브러리를 사용하여 특정 클래스에 대한 주석 및 이미지 정보를 추출하고 객체 감지에 사용할 수 있는 새로운 데이터 세트를 만들 수 있었습니다. 이 새로운 데이터 세트를 사용하여 특정 사용 사례에 맞게 미세 조정할 수 있는 사용자 지정 개체 감지 모델을 교육할 수 있습니다.