OpenCV Python Grabcut分割

【目标】

Grabcut 算法
创建一个交互程序

【理论】

从用户角度是如何工作的呢？用户在需要的目标上初始绘制一个矩形，前景目标必须完全在矩形内部，算法迭代的去分割然后得到更好的效果，但是有些情况下，分割效果不是很好，例如：会将部分前景标记为背景，反之亦然。这个时候，就需要用户做一些交互，告诉分割结果，哪些是前景哪些是背景，这样下次迭代就会得到更好的效果。

背景发生了哪些动作？

所有在矩形外的部分被标记为背景，在矩形框里面的都是未知。相似的所有用户给定的前景和背景都是 硬标签。
计算机不会初始化标签，主要取决于用户给定的数据，标记前景和背景像素。
GMM（混合高斯模型）被用户前景和背景
根据给定的数据，GMM学习和创造一些新的像素分布，位置像素会被标记为可能前景或背景，取决于其他硬标签像素（有点类似聚类）
图通过像素分布创建，图中的节点是像素，增加两个节点，一个是源节点，一个是汇节点，每个前景像素与源节点连接，每个背景像素与汇节点连接。
连接像素直接边的权重主要根据像素属于前景或背景的概率而定，像素之间的权重根据边缘信息和像素相似度而定。如果像素颜色有很大的不同，那么他们之间的边缘权重会小一些。
mincut算法用于分割图，它以最小代价函数将图分成两个分离的源节点和汇节点。代价函数是所有边权重之和。cut之后，所有连接到源节点的像素变成了前景，连接到汇节点的像素变成了背景。
这个过程持续到分类收敛。

在这里插入图片描述

【代码】

在这里插入图片描述

import numpy as np 
import cv2 

import sys 


COLOR_BLUE = [255, 0, 0]      # 矩形框颜色
COLOR_RED  = [0, 0, 255]      # 可能背景绘制颜色
COLOR_GREEN = [0, 255, 0]     # 可能前景绘制颜色
COLOR_BLACK = [0, 0, 0]       # 背景绘制颜色
COLOR_WHITE = [255, 255, 255] # 前景绘制颜色

DRAW_BG = {'color': COLOR_BLACK, 'val': 0}    # 背景，标记为0
DRAW_FG = {'color': COLOR_WHITE, 'val': 1}    # 前景，标记为1
DRAW_PR_BG = {'color': COLOR_RED, 'val': 2}    # 可能背景，标记为2
DRAW_PR_FG = {'color': COLOR_GREEN, 'val': 3}  # 可能前景，标记为3


class GrabCutApp():
    """
    GrabCutApp 利用grabcut对图像进行前景提取

    USAGE:
        python grabcut.py <image_name>

    README FIRST:
        two windows will show up, one for input and one for output 
      
        at first, in input window, right mouse buttom draw a rectangle around the object with blue. 
        then, press 'n' to segment the object (once of a few times). if the effect is not so good, 
        please prees key below what you want, the key concept is below:
      
        key '0' - select area of sure background
        key '1' - select area of sure foreground
        key '2' - select area of probable background
        key '3' - select area of probable foreground
      
        key 'r' - reset the setup
        key 'n' - segment the object
        key 's' - save the segmented image to "image_name_grabcut.png"
        key 'q' - quit
        key esc - quit
    """

    # 初始化
    def __init__(self, imagename: str) -> None:
        self.img = cv2.imread(imagename)
        if self.img is None:
            print('图像读取失败')
            sys.exit(0)
          
        self.rect = (0, 0, 1, 1) # 矩形框初始化
        self.drawing = False
        self.rectangle = False # 是否开始绘制矩形框
        self.rect_over = False # 判断矩形是否结束
        self.rect_or_mask = 100 # 矩形框或者mask的种类
        self.value = DRAW_FG
        self.thickness = 3
        self.radius = 5
      
      
    # 鼠标回调
    def onmouse(self, event, x, y, flags, param) -> None:
        # 自定义鼠标回调函数
        if event == cv2.EVENT_RBUTTONDOWN:
            self.rectangle = True
            self.ix, self.iy = x, y
        elif event == cv2.EVENT_MOUSEMOVE:
            if self.rectangle == True:
                self.img = self.img2.copy()
                cv2.rectangle(self.img, (self.ix, self.iy), (x, y), 
                            COLOR_BLUE, self.thickness)
                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 == cv2.EVENT_RBUTTONUP:
            self.rectangle = False
            self.rect_over = True
            cv.rectangle(self.img, (self.ix, self.iy), (x, y), 
                        COLOR_BLUE, self.thickness)
            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 == cv2.EVENT_LBUTTONDOWN:
            if self.rect_over == False:
                print("draw object first \n")
            else:
                self.drawing = True
                cv2.circle(self.img, (x, y), self.radius, self.value['color'], -1)
                cv2.circle(self.mask, (x, y), self.radius, self.value['val'], -1)
        elif event == cv2.EVENT_MOUSEMOVE:
            if self.drawing == True:
                cv2.circle(self.img, (x, y), self.radius, self.value['color'], -1)
                cv2.circle(self.mask, (x, y), self.radius, self.value['val'], -1)
        elif event == cv2.EVENT_LBUTTONUP:
            if self.drawing == True:
                self.drawing = False
                cv2.circle(self.img, (x, y), self.radius, self.value['color'], -1)
                cv2.circle(self.mask, (x, y), self.radius, self.value['val'], -1)

    def run(self):
        # 拷贝
        self.img2 = self.img.copy()
        # 初始化一个mask图像
        self.mask = np.zeros(self.img.shape[:2], dtype=np.uint8)
        self.output = np.zeros(self.img.shape, np.uint8)

        cv2.namedWindow('output')
        cv2.namedWindow('input')
        cv2.setMouseCallback('input', self.onmouse)
        cv2.moveWindow('input', self.img.shape[1] + 10, 0)

        print('draw a rectangle around the object use right mouse button to draw')

        while(1):
            cv2.imshow('output', self.output)
            cv2.imshow('input', self.img)
            k = cv2.waitKey(1)
          
            # 
            if k == 27 or k == ord('q'): # esc or 'q' to quit 
                break
            elif k == ord('0'): # BG drawing
                print("mark background regions with left mouse buttom \n")
                self.value = DRAW_BG
            elif k == ord('1'): # FG drawing
                print("mark foreground regions with left mouse buttom \n")
                self.value = DRAW_FG
            elif k == ord('2'): # PR_BG drawing
                self.value = DRAW_PR_BG
            elif k == ord('3'): # PR_FG drawing
                self.value = DRAW_PR_FG
            elif k == ord('s'): # save image
                bar = np.zeros((self.img.shape[0], 5, 3), np.uint8)
                res = np.hstack((self.img2, bar, self.img, bar, self.output))
                cv2.imwrite('grabcut_output_result.png', res)
                print('result saved as grabcut_output_result.png\n')
            elif k == ord('r'): # restore to original status
                print('reset all settings ...\n')
                self.rect = (0, 0, 1, 1)
                self.drawing = False
                self.rectangle = False 
                self.rect_or_mask = 100
                self.rect_over = False
                self.value = DRAW_FG
                self.img = self.img2.copy()

                self.mask = np.zeros(self.img.shape[:2], dtype=np.uint8)
                self.output = np.zeros(self.img.shape, np.uint8)
            elif k == ord('n'):
                print("for finer touchups, mark foreground and background after pressing keys 0-3")
                try:
                    bgdmodel = np.zeros((1, 65), np.float64)
                    fgdmodel = np.zeros((1, 65), np.float64)
                    if (self.rect_or_mask == 0): # original rectangle
                        cv2.grabCut(self.img2, self.mask, self.rect, bgdmodel, 
                                    fgdmodel, 1, cv2.GC_INIT_WITH_RECT)
                        self.rect_or_mask = 1
                    elif (self.rect_or_mask == 1): # grabcut after rectangle
                        cv2.grabCut(self.img2, self.mask, self.rect, bgdmodel, 
                                    fgdmodel, 5, cv2.GC_INIT_WITH_MASK)
                except:
                    import traceback
                    traceback.print_exc()
            mask2 = np.where((self.mask==1)+ (self.mask==3), 255, 0).astype('uint8')
            cv2.imshow('mask2', mask2)
            self.output = cv2.bitwise_and(self.img2, self.img2, mask=mask2)
      
        cv2.destroyAllWindows()


if __name__ == '__main__':
    app = GrabCutApp('assets/messi5.jpg')
    print(app.__doc__)
    app.run()

【接口】

cv.grabCut(	img, mask, rect, bgdModel, fgdModel, iterCount[, mode]	) ->	mask, bgdModel, fgdModel

执行 grabcut 算法

img: 输入8位3通道图像
mask: 输入输出的8位单通道图像，用矩形初始化
rect: ROI矩形，在矩形外面的部分被认为是背景，只有当 mode = GC_INIT_WITH_RECT 时才有效
bgdModel: 存储背景模型的参数，处理同一个图像时，不要修改该模型
fgdModel: 存储前景模型的参数，处理同一个图像时，不要修改该模型
iterCount: 迭代次数
mode: 不同的模式 GrabCutModes