【OpenCV-Python】教程：3-16 利用Grabcut交互式前景提取

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 

    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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178

【接口】

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

执行 grabcut 算法

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

• GrabCutModes

【参考】

1. OpenCV 官方文档
2. "GrabCut": interactive foreground extraction using iterated graph cuts
3. GrabCut image segmentation algorithm.