OpenCV使用 GrabCut 算法进行交互式前景提取

Posted 2022-02-02 程序媛一枚~

OpenCV使用 GrabCut 算法进行交互式前景提取

• • 1. 效果图
  • 2. 源码
  • 参考

这篇博客将介绍如何使用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