OpenCV使用 GrabCut 算法进行交互式前景提取
这篇博客将介绍如何使用Python,OpenCV中的GrabCut 算法来提取图像中的前景,并为此创建一个交互式应用程序。
1. 效果图
官方示例——lena 原始图 VS grabcut前景设置 VS 前景抠图效果如下:
鼠标右键绘制矩形,按下 ‘0’ 左键点击标记确定背景,按下 ‘1’ 左键标记确定前景,按 ‘r’ 键重置矩形等绘制,按 ’s‘ 键保存结果图,然后得到比较干净友好的效果图; 我最喜欢的颖宝,原始图 VS GrabCut图 VS 抠图效果如下:
2. 源码
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}
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):
if len(sys.argv) == 2:
filename = sys.argv[1]
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)
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:
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)
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
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):
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):
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()
参考
|