【OpenCV】GrabCutの使い方【セグメンテーション】

入社して半年が経ちました。最年少社員の安藤です。

つい一月半ほど前まで高校時代の同級生と同居していたのですが、最近のぼくのこころのなかにはコンピュータビジョン(CV)が住み着いています。
そんな彼はふとした時、ぼくに問いかけてきます。

「CVは楽しいかね？」

仕事は楽しいかね？ (きこ書房)

作者: デイル・ドーテン
出版社/メーカー: きこ書房　
発売日: 2015/01/19
メディア: Kindle版
この商品を含むブログ (6件) を見る

楽しいです。

そんなわけで今回はOpenCVの便利なメソッドのひとつを紹介したいと思います。

GrabCutとは？

GrabCutとはセグメンテーションに利用されるアルゴリズムの名称です。

グラフ理論を応用したコスト最小化手法GraphCutを用いたくりかえし処理を行っているそうです。
アルゴリズム本体の詳細については他にまとめが存在しているので、下記をご覧ください。

http://visitlab.jp/pdf/GrabCut.pdf
http://zellij.hatenablog.com/entry/20131004/p1

むずかしそうな数式がたくさん並んでいますね。

こんな時もぼくの精神的CVはいつもそばにいて、助言をあたえてくれるのです。

「OpenCVを使え」

OpenCVを使いましょう。

解説

import cv2 as cv
import numpy as np


BLUE = [255,0,0]        # rectangle color
RED = [0,0,255]         # possible BG
GREEN = [0,255,0]       # possible FG
BLACK = [0,0,0]         # obvious BG
WHITE = [255,255,255]   # obvious FG

DRAW_BG = {'color' : BLACK, 'val' : 0}
DRAW_FG = {'color' : WHITE, 'val' : 1}
DRAW_PR_FG = {'color' : GREEN, 'val' : 3}
DRAW_PR_BG = {'color' : RED, 'val' : 2}


# 現在の描画する線の色
value = DRAW_FG


# 描画する線の太さ
thickness = 3


# ウィンドウ名
win_name = "WINDOW"


# grabcut結果表示の際、alphaが大きいほど黒背景が濃ゆくなる
alpha = 0.7
beta = 1 - alpha


# grabcutイテレーション回数
itr = 1


# 描画フラグ
drawing = False


# 出力イメージ番号
pic_num = 1



def mouse_event(event, x, y, flags, param):
    global display, mask, thickness, value, drawing, bgdModel, fgdModel

    if event == cv.EVENT_LBUTTONDOWN:
        drawing = True
        cv.circle(display, (x, y), thickness, value['color'], -1)
        cv.circle(mask, (x, y), thickness, value['val'], -1)

        
    elif event == cv.EVENT_MOUSEMOVE:
        if drawing == True:
            cv.circle(display, (x, y), thickness, value['color'], -1)
            cv.circle(mask, (x, y), thickness, value['val'], -1)
        
        
    elif event == cv.EVENT_LBUTTONUP:
        if drawing == True:
            drawing = False
            cv.circle(display, (x, y), thickness, value['color'], -1)
            cv.circle(mask, (x, y), thickness, value['val'], -1)
            

if __name__ == '__main__':


    img = cv.imread("lena.png")
    display = img.copy()

    mask = np.zeros(img.shape[:2], dtype=np.uint8) 
    bgdModel = np.zeros((1,65), dtype=np.float64)
    fgdModel = np.zeros((1,65), dtype=np.float64)

    cv.namedWindow(win_name, cv.WINDOW_NORMAL)

    roi = cv.selectROI(win_name, img)

    cv.grabCut(img, mask, roi, bgdModel, fgdModel, itr, cv.GC_INIT_WITH_RECT)
    
    mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
    display = img * mask2[:,:,np.newaxis]
    display = cv.addWeighted(display, alpha, img, beta, 0)

    cv.setMouseCallback(win_name, mouse_event)
    
    while (True):

        cv.imshow(win_name, display)
        key = cv.waitKey(1)

        if key == 27:
            cv.destroyAllWindows()
            break


        elif key == ord("n"):

            drawing = False

            cv.grabCut(img, mask, roi, bgdModel, fgdModel, itr, cv.GC_INIT_WITH_MASK)
            mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
            display = img * mask2[:,:,np.newaxis]
            display = cv.addWeighted(display, alpha, img, beta, 0)


        elif key == ord("s"):
            cv.imwrite("processed" + str(pic_num) + ".png", display)
            pic_num += 1


        elif key == ord('0'): # BG drawing
            print(" mark background regions with left mouse button \n")
            value = DRAW_BG
            print("Obvious BGD")


        elif key == ord('1'): # FG drawing
            print(" mark foreground regions with left mouse button \n")
            value = DRAW_FG
            print("Obvious FGD")


        elif key == ord('2'): # PR BG drawing
            value = DRAW_PR_BG
            print("Possible BGD")


        elif key == ord('3'): # PR FG drawing
            value = DRAW_PR_FG
            print("Possible FGD")