OpenCV使用 GrabCut 演算法進行互動式前景提取

這篇博客將介紹如何使用Python，OpenCV中的GrabCut 演算法來提取影像中的前景，并為此創建一個互動式應用程式，

1. 效果圖

官方示例——lena 原始圖 VS grabcut前景設定 VS 前景摳圖效果如下：

滑鼠右鍵繪制矩形，按下 ‘0’ 左鍵點擊標記確定背景，按下 ‘1’ 左鍵標記確定前景，按 ‘r’ 鍵重置矩形等繪制，按 ’s‘ 鍵保存結果圖，然后得到比較干凈友好的效果圖；


我最喜歡的穎寶，原始圖 VS GrabCut圖 VS 摳圖效果如下：

2. 原始碼

# 使用GrabCut進行互動式的前景提取（Interactive Image Segmentation using GrabCut algorithm）
# USAGE:
# python grabcut.py ml2.jpg

# 運行會呈現倆個視窗，輸入影像及輸出影像視窗
# 首先用滑鼠右鍵繪制一個前景的矩形框，按下 'n' 鍵進行前景提取
# 然后按下滑鼠左鍵繪制任何區域，點、線等不規則區域，按下 'n' 鍵進行前景提取

# 按下鍵 '0' - 選擇確定背景
# 按下鍵 '1' - 選擇確定前景
# 按下鍵 '2' - 選擇可能背景
# 按下鍵 '3' - 選擇可能前景
# 按下鍵 'n' - 更新分割結果
# 按下鍵 'r' - 重置分割
# 按下鍵 's' - 保存分割的結果
# 按下鍵 'ESC' - 退出分割

from __future__ import print_function

import sys

import cv2 as cv
import numpy as np


class App():
    BLUE = [255, 0, 0]  # 矩形顏色
    RED = [0, 0, 255]  # 可能背景
    GREEN = [0, 255, 0]  # 可能前景
    BLACK = [0, 0, 0]  # 確定背景
    WHITE = [255, 255, 255]  # 確定前景

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

    # 設定flags
    rect = (0, 0, 1, 1)
    drawing = False  # 是否繪制曲線
    rectangle = False  # 是否繪制矩形
    rect_over = False  # 是否矩形繪制完畢
    rect_or_mask = 100  # 選擇矩形或者蒙版
    value = DRAW_FG  # 繪制初始化的前景
    thickness = 3  # 筆刷粗細

    def onmouse(self, event, x, y, flags, param):
        # 繪制矩形（按下滑鼠右鍵）
        if event == cv.EVENT_RBUTTONDOWN:
            self.rectangle = True
            self.ix, self.iy = x, y

        elif event == cv.EVENT_MOUSEMOVE:
            if self.rectangle == True:
                self.img = self.img2.copy()
                cv.rectangle(self.img, (self.ix, self.iy), (x, y), self.BLUE, 2)
                self.rect = (min(self.ix, x), min(self.iy, y), abs(self.ix - x), abs(self.iy - y))
                self.rect_or_mask = 0

        elif event == cv.EVENT_RBUTTONUP:
            self.rectangle = False
            self.rect_over = True
            cv.rectangle(self.img, (self.ix, self.iy), (x, y), self.BLUE, 2)
            self.rect = (min(self.ix, x), min(self.iy, y), abs(self.ix - x), abs(self.iy - y))
            self.rect_or_mask = 0
            print(" Now press the key 'n' a few times until no further change \n")

        # 繪制接觸曲線
        if event == cv.EVENT_LBUTTONDOWN:
            if self.rect_over == False:
                print("first draw rectangle \n")
            else:
                self.drawing = True
                cv.circle(self.img, (x, y), self.thickness, self.value['color'], -1)
                cv.circle(self.mask, (x, y), self.thickness, self.value['val'], -1)

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

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

    def run(self):
        # Loading images
        if len(sys.argv) == 2:
            filename = sys.argv[1]  # for drawing purposes
        else:
            print("No input image given, so loading default image, ml.jpg \n")
            print("Correct Usage: python grabcut.py <filename> \n")
            filename = 'ml.jpg'

        print(filename)
        self.img = cv.imread(filename)
        self.img2 = self.img.copy()  # 復制原始影像
        self.mask = np.zeros(self.img.shape[:2], dtype=np.uint8)  # mask初始化為PR_BG可能背景
        self.output = np.zeros(self.img.shape, np.uint8)  # 輸出影像

        # 輸入輸出框
        cv.namedWindow('output')
        cv.namedWindow('input')
        cv.setMouseCallback('input', self.onmouse)
        cv.moveWindow('input', self.img.shape[1] + 10, 90)

        print(" Instructions: \n")
        print(" Draw a rectangle around the object using right mouse button \n")

        while (1):

            cv.imshow('output', self.output)
            cv.imshow('input', self.img)
            k = cv.waitKey(1)

            # 按鍵事件系結
            if k == 27:  # 按ESC退出
                break
            elif k == ord('0'):  # 繪制背景
                print(" mark background regions with left mouse button \n")
                self.value = self.DRAW_BG
            elif k == ord('1'):  # 繪制前景
                print(" mark foreground regions with left mouse button \n")
                self.value = self.DRAW_FG
            elif k == ord('2'):  # 繪制可能背景
                self.value = self.DRAW_PR_BG
            elif k == ord('3'):  # 繪制可能前景
                self.value = self.DRAW_PR_FG
            elif k == ord('s'):  # 保存結果
                bar = np.zeros((self.img.shape[0], 10, 3), np.uint8)
                res = np.hstack((self.img2, bar, self.img, bar, self.output))
                cv.imwrite('ml_grabcut_output.jpg', res)
                print(" Result saved as image \n")
            elif k == ord('r'):  # 重置，重新進行選中前景背景
                print("resetting \n")
                self.rect = (0, 0, 1, 1)
                self.drawing = False
                self.rectangle = False
                self.rect_or_mask = 100
                self.rect_over = False
                self.value = self.DRAW_FG
                self.img = self.img2.copy()
                self.mask = np.zeros(self.img.shape[:2], dtype=np.uint8)  # mask initialized to PR_BG
                self.output = np.zeros(self.img.shape, np.uint8)  # output image to be shown
            elif k == ord('n'):  # 分離圖片
                print(""" For finer touchups, mark foreground and background after pressing keys 0-3
                and again press 'n' \n""")
                try:
                    bgdmodel = np.zeros((1, 65), np.float64)
                    fgdmodel = np.zeros((1, 65), np.float64)
                    if (self.rect_or_mask == 0):  # 使用矩形進行grabcut
                        cv.grabCut(self.img2, self.mask, self.rect, bgdmodel, fgdmodel, 1, cv.GC_INIT_WITH_RECT)
                        self.rect_or_mask = 1
                    elif (self.rect_or_mask == 1):  # 使用mask進行grabcut
                        cv.grabCut(self.img2, self.mask, self.rect, bgdmodel, fgdmodel, 1, cv.GC_INIT_WITH_MASK)
                except:
                    import traceback
                    traceback.print_exc()

            mask2 = np.where((self.mask == 1) + (self.mask == 3), 255, 0).astype('uint8')
            self.output = cv.bitwise_and(self.img2, self.img2, mask=mask2)

        print('Done')


if __name__ == '__main__':
    print(__doc__)
    App().run()
    cv.destroyAllWindows()

參考

• https://docs.opencv.org/3.0-beta/doc/py_tutorials/py_imgproc/py_grabcut/py_grabcut.html#grabcut