COCO dataset 다루기

※ 본 글은 <딥러닝 컴퓨터비전 완벽가이드> 강의 및 여러 자료를 참고하여 쓰여진 글입니다. ※

github: github.com/chaeyeongyoon/ComputerVision_Study

chaeyeongyoon/ComputerVision_Study

pycocotools API를 이용해 COCO dataset을 다뤄보겠습니다.

1. 이미지의 segmentation 시각화하기

1. annotation file(.json)을 COCO객체로 로드 -> instances 대한 파일로
2. 내가 찾고 싶은 객체들의 id리스트 로드 coco.getIds(catnms=[ ])
3. 그 객체들을 가지고 있는 이미지의 id 리스트 로드 coco.getImgIds(catIds=[ ])
4. 그 중 한 이미지의 정보 로드(coco.load_Imgs(img_id)[0])
5. 이미지 정보의 file_name 이용해 경로를 지정해 plt로 이미지 띄워준 후
6. annotation id들 list로드(getAnnIds())
7. annotatioin id로 coco.loadAnns(annIds) : annotation 정보들 로드
8. 이미지 plt로 로드 한후 
9. coco.showAnns( )로 segmentation 로드

from pycocotools.coco import COCO
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

data_dir = 'D:/jupyter_dir/DLCV/data/coco'
data_type = 'val2017'
annFile = '{}/annotations_trainval2017/annotations/instances_{}.json'.format(data_dir, data_type)
coco = COCO(annFile) # COCO객체로 만들어줌

catIds = coco.getCatIds(catNms=['person', 'dog', 'skateboard'])
imgIds = coco.getImgIds(catIds=catIds) # [549220, 324158, 279278]

img = coco.loadImgs(324158)
'''
[{'license': 1, 'file_name': '000000324158.jpg', 'coco_url': 'http://images.cocodataset.org/val2017/000000324158.jpg', 'height': 334, 'width': 500, 'date_captured': '2013-11-19 23:54:06', 'flickr_url': 'http://farm1.staticflickr.com/169/417836491_5bf8762150_z.jpg', 'id': 324158}]
'''
img = coco.loadImgs(324158)[0]

annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
anns = coco.loadAnns(annIds)

plt.figure(figsize=(8, 10))
plt.imshow(img_array)
plt.axis('off')
coco.showAnns(anns)

2. Bounding Box 시각화하기

앞에서 구한 anns 사용

anns = [{'segmentation': [[216.7, 211.89, 216.16, 217.81, 215.89, 220.77, 215.89, 223.73, 217.77, 225.35, 219.12, 224.54, 219.12, 220.5, 219.66, 217.27, 219.93, 212.7, 220.46, 207.85, 219.66, 203.01, 218.85, 198.43, 217.77, 195.74, 216.7, 194.93, 215.62, 190.62, 215.62, 186.59, 214.27, 183.89, 211.85, 184.16, 211.85, 187.66, 210.24, 187.66, 209.16, 184.97, 207.81, 183.36, 205.12, 186.59, 205.12, 189.28, 201.08, 192.78, 199.74, 195.2, 196.78, 200.04, 196.51, 203.01, 198.12, 205.43, 197.32, 209.2, 196.78, 213.23, 197.05, 218.89, 199.74, 221.85, 201.62, 225.35, 201.62, 233.69, 201.08, 236.11, 202.97, 236.38, 204.85, 236.11, 204.58, 232.34, 203.78, 228.85, 205.39, 233.15, 207.81, 235.57, 208.62, 234.23, 206.74, 231.27, 205.12, 228.04, 206.74, 222.39, 208.35, 219.96, 210.77, 217.54, 211.85, 221.85, 214.54, 223.73, 212.93, 217.54, 212.93, 215.66, 215.89, 212.96, 216.16, 212.16]], 'area': 759.3375500000002, 'iscrowd': 0, 'image_id': 324158, 'bbox': [196.51, 183.36, 23.95, 53.02], 'category_id': 18, 'id': 10673}, ......(생략).... ]

 

for ann in anns: 반복문 안에서 ann['bbox']와 ann['category_id']를 이용해 시각화를 해줍니다.

image_array = cv2.imread(img_file_path)
image_array = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)

for ann in anns:
    ann_bbox = list(map(int, ann['bbox']))
    ann_category_id = ann['category_id']
    
    pose = tuple(ann_bbox)
    cv2.rectangle(image_array, pose, (125, 255, 51), thickness=2)
    
    category = coco.loadCats(ann_category_id)[0]
    category_name = category['name']
    cv2.putText(image_array, category_name, (ann_bbox[0]-5, ann_bbox[1]-5),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (125, 255, 51), 1)
    
plt.figure(figsize=(8, 10))
plt.imshow(image_array)
plt.axis('off')