cv2 Grabcut图像分割

目录

1、概述

2、算法流程

3、函数

4、示例代码

---

1、概述

Grabcut是基于图割(graph cut)实现的图像分割算法，它需要用户输入一个bounding box作为分割目标位置，实现对目标与背景的分离/分割。Grabcut分割速度快，效果好，支持交互操作，因此在很多APP图像分割/背景虚化的软件中经常使用。

2、算法流程

• 在图片中定义含有（一个或多个）物体的矩形；
• 矩形外的区域被自动认为是背景；
• 对于用户定义的矩形区域，可用背景中数据来区分是前景还是背景；
• 用高斯混合模型（GMM）来对背景和前景建模，并将未定义的像素标记为可能的前景或背景；
• 图像中的每一个像素都被看作通过通过虚拟边与周围像素连接，而每条边都有一个属于前景或背景的概率这基于它和周围像素颜色上的相似性；
• 每一个像素（即算法中的节点）会与前一各前景或背景节点连接；
• 在节点连接完成后，用图论中最大流最小割的方法来分割。

3、函数

cv2.grabCut(img, rect, mask, bgdModel, fgdModel, iterCount, mode = GC_EVAL)

其中：

img --> 输入的三通道图像；

mask --> 输入的单通道图像，初始化方式为GC_INIT_WITH_RECT表示ROI区域可以被初始化为：

• GC_BGD --> 定义为明显的背景像素 0

• GC_FGD --> 定义为明显的前景像素 1

• GC_PR_BGD --> 定义为可能的背景像素 2

• GC_PR_FGD --> 定义为可能的前景像素 3

• rect --> 表示roi区域；

• bgdModel --> 表示临时背景模型数组；

• fgdModel --> 表示临时前景模型数组；

• iterCount --> 表示图割算法迭代次数, 次数越多，效果越好；

• mode --> 当使用用户提供的roi时候使用GC_INIT_WITH_RECT。

4、示例代码

输入：

采用 selectROI, 可以在图中自己选定ROI区域：

选定后，按enter 或则 Space 进行grabcut；

重新选ROI，只需用鼠标重新选择即可；

按 c 结束程序。

1）读取图像，使用cv2选择矩形框

import cv2
import numpy as np

img = cv2.imread(img_path)
r = cv2.selectROI('input', img, False)  # 返回 (x_min, y_min, w, h)
print("input:", r)

按回车，输出：

input: (100, 65, 417, 473)

2）读取ROI，执行grabcut

# roi区域
roi = img[int(r[1]):int(r[1] + r[3]), int(r[0]):int(r[0] + r[2])]

# 原图mask
mask = np.zeros(img.shape[:2], dtype=np.uint8)

# 矩形roi
rect = (int(r[0]), int(r[1]), int(r[2]), int(r[3]))  # 包括前景的矩形，格式为(x,y,w,h)

bgdmodel = np.zeros((1, 65), np.float64)  # bg模型的临时数组
fgdmodel = np.zeros((1, 65), np.float64)  # fg模型的临时数组

cv2.grabCut(img, mask, rect, bgdmodel, fgdmodel, 11, mode=cv2.GC_INIT_WITH_RECT)

# 提取前景和可能的前景区域
mask2 = np.where((mask == 1) + (mask == 3), 255, 0).astype('uint8')
print(mask2.shape)

result = cv2.bitwise_and(img, img, mask=mask2)
cv2.imwrite('result.jpg', result)
cv2.imwrite('roi.jpg', roi)

cv2.imshow('roi', roi)
cv2.imshow("result", result)
cv2.waitKey(0)
cv2.destroyAllWindows()

输出结果：