Pythonでポリゴン(多角形)の領域を抽出する方法

0. 概要

COCO Formatを使っているとき、ある物体のみを抽出したいなんていうことがよくある。
その際、ポリゴン情報をマスクにして画像から物体を抽出するという処理が必要である。
この方法が探しても意外に見つからない。今回はそんな方法を紹介する。

1. コード

一先ず、以下のコードを見れば秒で分かるはずである。

import cv2

src_img = cv2.imread("YOUR_PATH")

# parameter
sg = [1, 1, 2, 2, 3, 3]
bb = [1, 1, 2, 2]

# working file
mask = np.zeros_like(src_img) # (y, x, c)

# segmentation data
sg = np.asarray(sg)
poly_number = int(len(sg)/2)
poly = np.zeros( (poly_number, 2) )
for i in range(poly_number):
  poly[i][0] = sg[(i * 2) + 0] # x
  poly[i][1] = sg[(i * 2) + 1] # y

# generate mask
mask = cv2.fillConvexPoly(mask, np.array(poly, 'int32'), color=(255, 255, 255))

# generate src_img and mask_image
cv2.imwrite("./src_img.png", src_img)
cv2.imwrite("./mask.png", mask)

# masked image
masked_src_img = np.where(mask==255, src_img, mask)
cv2.imwrite("./procData/masked.png", masked_src_img)

# cropping img
bb = np.asarray(bb, 'int32')
offset_x = bb[0]
offset_y = bb[1]
length_x = bb[2]
length_y = bb[3]
cv2.imwrite("./masked_crop.png", masked_src_img[offset_y : (offset_y + length_y), offset_x : (offset_x + length_x)])

まず、以下でポリゴンからマスク画像へ変換している。
注意すべきは、int32以外エラーが発生することである。

cv2.fillConvexPoly(mask, np.array(poly, 'int32'), color=(255, 255, 255))

次に、マスク画像が255値の部分はsrc_imgの値を使うと指定すればよいだけである。

masked_src_img = np.where(mask==255, src_img, mask)

最後にBounding boxでCropingすれば特定領域の抽出ができる。

Enjoy !