Segmentation Masking 시각화

pycocotools는 coco dataset의 이미지나 annotation등의 정보를 쉽게 접근할 수 있게 지원해주는 API이다.

 

from pycocotools.coco import COCO
import numpy as np

pycocotools는 이미 colab에 설치가 되어있기 때문에 별도의 설치과정을 거치지 않아도 된다.

 

사용법은 간단하다.

 왼쪽과 같은 그림에서 coco dataset id에 속하는 person, dog, skateboard를 추출하고 싶은 object로 설정한 뒤 pycocotools를 이용하여 segmentation을 진행하면된다.

 

 

 

 

 

 

 

catIds = coco.getCatIds(catNms=['person','dog','skateboard']);
annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
anns = coco.loadAnns(annIds)
plt.figure(figsize=(12, 14))
plt.imshow(image_array)
plt.axis('off')

coco.showAnns(anns)

getAnnIds()로 특정 이미지에 속하는 annotation id 정보를 가져온 후에 이 id정보를 loadAnns()의 입력으로 하여 해당 이미지의 모든 annotation 정보를 가져온다. 이때의 segmentation정보는 polygon형태의 구조로 되어있다.

이후 coco.showAnns()에 받아온 annotation정보를 입력으로 하면 segmentation이 시각화까지 완료된다.

하지만 이렇게 짜여져 있는 라이브러리를 사용하지 못하는 경우도 있기 때문에 polygon형태의 구조를 opencv를 활용해서 직접 segmentation을 수행해줄 수도 있다.

---

앞에서 추출한 annotation정보를 가지는 anns의 인덱스2번째 값을 활용하면 polygon 형태의 segmentation에 필요한 정보를 활용할 수 있다.

ann_2_seg = anns[2]['segmentation'][0]

polygon_x = [x for index, x in enumerate(ann_2_seg) if index % 2 == 0]
polygon_y = [x for index, x in enumerate(ann_2_seg) if index % 2 == 1]

polygon_xy = [[x, y] for x, y in zip(polygon_x, polygon_y)]

anns의 2번째 인덱스 값중 첫번째 값이 img에서 person에 대한 annotation정보를 가지고 있다.

이 정보를 활용하여 x좌표 y좌표를 각각 추출하여 polygon_xy라는 값을 만들어준다.

import numpy as np

green_color = (0, 255, 0)

draw_img = image_array.copy()
polygon_xy = np.array(polygon_xy, np.int32)
draw_img = cv2.polylines(draw_img, [polygon_xy], True, (0, 255, 0))

plt.figure(figsize=(12, 14))
plt.imshow(draw_img)
plt.axis('off')

이후 opencv의 polylines함수를 이용하여 person에 대한 외곽선을 그려줄 수 있다.

green_color = (0, 255, 0)

draw_img = image_array.copy()
polygon_xy = np.array(polygon_xy, np.int32)
draw_img = cv2.fillPoly(draw_img, [polygon_xy], (0, 255, 0))

plt.figure(figsize=(12, 14))
plt.imshow(draw_img)
plt.axis('off')

opencv의 fillpoly함수를 사용하면 instance segmentation을 수행할 수 있다.