OpenCV uses GrabCut algorithm for interactive foreground extraction
This blog will show how to use Python, GrabCut algorithm in OpenCV to extract foreground in an image and create an interactive application for this.
1. Rendering
The official example - lena original image VS grabcut foreground setting VS foreground matting effect is as follows:
Use the right mouse button to draw a rectangle, press the '0' left button to click the marker to determine the background, press the '1' left button to determine the foreground, press the 'r' key to reset the rectangle, etc., press the 's' key to save the resulting graph, and then get the comparison Clean and friendly renderings;
my favorite Yingbao, the original image VS GrabCut image VS cutout effect is as follows:
2. Source code
# 使用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()