[Dataset] Pycocotools를 이용한 COCO Dataset 확인하기

COCO Dataset

이번 포스팅은 Pycocotools를 이용하여 coco dataset에 데이터들이 어떤 형태로 있는지 확인합니다. 

데이터 셋은 아래 링크에서 다운 받을 수 있으며

저는  /hostDir/dataSet/images  와  /hostDir/dataSet/annotation  폴더에 다운받아 코드를 실행합니다.

https://cocodataset.org/#download

COCO - Common Objects in Context

 

 

---

1. COCO 초기화

COCO 데이터 셋은 사진들과 json 파일로 구성되어 있습니다.

우선 필요로하는 모듈들을 import하고 json 파일 경로를 COCO 객체의 입력으로 넣어 coco를 생성합니다.

import os
import torch
import matplotlib.pyplot as plt

from pycocotools.coco import COCO
from PIL import Image

dataDir = '/hostDir/dataSet'
annVal = os.path.join(dataDir, 'annotations', 'person_keypoints_val2014.json')

coco = COCO(annVal)
print(coco)

출력:

loading annotations into memory...
Done (t=2.65s)
creating index...
index created!
<pycocotools.coco.COCO object at 0x7f28c468c290>

 

 

---

2. 카테고리 인덱스 정보

pycocotools에서는 COCO 데이터 셋을 활용하기 위한 여러 유용한 함수들을 제공하고 있습니다. 

오늘 사용할 함수 및 자주 사용하는 함수들은 아래와 같습니다.

• getAnnIds: Get ann ids that satisfy given filter conditions.
• getCatIds: Get cat ids that satisfy given filter conditions.
• getImgIds: Get img ids that satisfy given filter conditions.
• loadAnns: Load anns with the specified ids.
• loadCats: Load cats with the specified ids.
• loadImgs: Load imgs with the specified ids.
• showAnns: Display the specified annotations.

해당 내용은 아래 github에서 확인할 수 있고 어떤한 역할을 하는지도 코드를 통해 자세히 알 수 있습니다.

https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/coco.py

GitHub - cocodataset/cocoapi: COCO API - Dataset @ http://cocodataset.org/

 

 

우선 카테고리 인덱스 정보를 getCatIds를 이용하여 획득합니다.

아래와 같이 원하는 카테고리를 입력으로 넣으면 해당 인덱스를 리턴합니다.

 

참고로, 주석 처리된 부분은 함수 정의로 어떤 입력 값을 받는지, default 값이 무엇인지 확인하기 위해 작성하였습니다.

# def getCatIds(self, catNms=[], supNms=[], catIds=[]):
personIds = coco.getCatIds(catNms=['person']) 
print(f'person index: \n{personIds}\n')

출력:

person index: 
[1]

참고로 카테고리의 종류를 알고 싶다면 loadCats을 이용하여 확인할 수 있습니다.

저는  person_keypoints_val2014.json 을 사용하여 사실상 카테고리가 person 밖에 없습니다.

 

 

 

---

3. 카테고리 인덱스를 이용한 영상 추출

앞서 person 카테고리의 인덱스를 획득하였습니다. 

획득한 인덱스를 이용하여 카테고리가 포함되어 있는, 즉 person이 포함되어 있는 영상들만 추출해 보겠습니다.

 

우선 getImgIds에 앞서 구한 인덱스를 입력하여 해당 카테고리가 포함되어 있는 모든 영상 인덱스를 획득합니다.

그 중 하나의 영상을 선택합니다. (아래 코드에서는 101번째 영상)

getAnnIds를 통해 해당 사진의 Annotation 인덱스를 얻고, 이를 loadAnns에 입력하여 해당 영상의 Annotation을 획득할 수 있습니다.

# def getImgIds(self, imgIds=[], catIds=[]):
ids = coco.getImgIds(catIds = personIds)

imgIds = ids[100]
# def getAnnIds(self, imgIds=[], catIds=[], areaRng=[], iscrowd=None):
annIds = coco.getAnnIds(imgIds = imgIds, catIds = personIds)
# def loadAnns(self, ids=[]):
anns = coco.loadAnns(annIds)

print(f"101th image's annotations: \n{anns}\n")

 

아래 출력에서 알 수 있듯이, Annotation에는 segmentation정보와 Pose estimation에 사용되는 keypoints, Object detection에 사용되는 bbox 등 정보를 확인 할 수 있습니다.

참고로 해당 사진에는 사람이 한 명으로 크기가 1인 리스트가 출력되었습니다. 여러 사람이 포함되어 있을 경우 사람 수 만큼의 정보가 포함됩니다.

출력:

101th image's annotations: 
[{'segmentation': [[46.07, 331.54, 64.04, 279.85, 84.27, 260.75, 98.88, 231.54, 108.99, 215.81,
119.1, 180.97, 139.33, 149.51, 157.3, 142.77, 183.15, 156.25, 193.26, 174.23, 196.63, 193.33, 
194.38, 216.93, 183.15, 237.15, 183.15, 259.63, 192.13, 277.6, 206.74, 268.61, 225.84, 266.37, 
235.96, 268.61, 249.44, 280.97, 259.55, 292.21, 248.31, 319.18, 231.46, 342.77, 214.61, 356.25, 
204.49, 374.23, 212.36, 424.79, 204.49, 460.75, 203.37, 487.72, 69.66, 487.72, 59.55, 400.07,
42.7, 346.14]], 'num_keypoints': 11, 'area': 45558.63015, 'iscrowd': 0, 'keypoints': [185, 227,
2, 188, 209, 2, 167, 210, 2, 0, 0, 0, 122, 206, 2, 153, 274, 2, 80, 298, 2, 186, 314, 2, 125,
432, 2, 220, 302, 2, 199, 346, 2, 161, 462, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
'image_id': 262642, 'bbox': [42.7, 142.77, 216.85, 344.95], 'category_id': 1, 'id': 443109}]

 

 

 

---

4. 영상 표시

위에서 구한 영상 인덱스(imgIds)를 loadImgs에 입력하면 해당 영상 정보를 얻을 수 있습니다.

imgInfo = coco.loadImgs(imgIds)
print(f'imgInfo: \n{imgInfo}\n')
print(f"fileName: \n{imgInfo[0]['file_name']}")

영상 정보에는 아래와 같이 이름인 file_name, 영상 크기 height와 width, 그리고 URL 정보인 coco_url 등이 포함되어 있습니다.

출력:

imgInfo: 
[{'license': 3, 'file_name': 'COCO_val2014_000000262642.jpg', 'coco_url': 
'http://images.cocodataset.org/val2014/COCO_val2014_000000262642.jpg', 
'height': 500, 'width': 375, 'date_captured': '2013-11-20 13:27:23', 
'flickr_url': 'http://farm1.staticflickr.com/179/423530889_a835ffd9cc_z.jpg', 
'id': 262642}]

fileName: 
COCO_val2014_000000262642.jpg

 

 

 

영상 정보에 있는 영상 이름(imgInfo[0]['file_name'])을 이용하여 해당 영상을 불러오고 표시해봅시다.

특히, showAnns를 이용하면 segmentation, keypoints 등을 표시할 수 있습니다.

imageName = os.path.join(dataDir, 'images', 'val2014', imgInfo[0]['file_name'])
image = Image.open(imageName).convert('RGB')

plt.imshow(image)
# def showAnns(self, anns, draw_bbox=False):
coco.showAnns(anns, draw_bbox = True)

출력:

 

 

감사합니다.