admin管理员组

文章数量:1122937

【OpenCV

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}  # 可能前景，标记为3class 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 backgroundkey '1' - select area of sure foregroundkey '2' - select area of probable backgroundkey '3' - select area of probable foregroundkey 'r' - reset the setupkey 'n' - segment the objectkey 's' - save the segmented image to "image_name_grabcut.png"key 'q' - quitkey 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 = Falseself.rectangle = False # 是否开始绘制矩形框self.rect_over = False # 判断矩形是否结束self.rect_or_mask = 100 # 矩形框或者mask的种类self.value = DRAW_FGself.thickness = 3self.radius = 5# 鼠标回调def onmouse(self, event, x, y, flags, param) -> None:# 自定义鼠标回调函数if event == cv2.EVENT_RBUTTONDOWN:self.rectangle = Trueself.ix, self.iy = x, yelif 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 = 0elif event == cv2.EVENT_RBUTTONUP:self.rectangle = Falseself.rect_over = Truecv.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 = 0print(" 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 = Truecv2.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 = Falsecv2.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 breakelif k == ord('0'): # BG drawingprint("mark background regions with left mouse buttom \n")self.value = DRAW_BGelif k == ord('1'): # FG drawingprint("mark foreground regions with left mouse buttom \n")self.value = DRAW_FGelif k == ord('2'): # PR_BG drawingself.value = DRAW_PR_BGelif k == ord('3'): # PR_FG drawingself.value = DRAW_PR_FGelif k == ord('s'): # save imagebar = 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 statusprint('reset all settings ...\n')self.rect = (0, 0, 1, 1)self.drawing = Falseself.rectangle = False self.rect_or_mask = 100self.rect_over = Falseself.value = DRAW_FGself.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 rectanglecv2.grabCut(self.img2, self.mask, self.rect, bgdmodel, fgdmodel, 1, cv2.GC_INIT_WITH_RECT)self.rect_or_mask = 1elif (self.rect_or_mask == 1): # grabcut after rectanglecv2.grabCut(self.img2, self.mask, self.rect, bgdmodel, fgdmodel, 5, cv2.GC_INIT_WITH_MASK)except:import tracebacktraceback.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

• GrabCutModes

【参考】

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

本文标签： opencv