概述

本篇博客主要介绍如何使用Mask R-CNN和 GrabCut进行自动抠图，主要分为以下部分

如何安装Mask R-CNN
如何使用Mask R-CNN

效果展示

如何安装Mask R-CNN

获取Mask R-CNN代码

你可以从GitHub下载Mask R-CNN的代码，自行集成。或者直接使用我编写的例子，我已经在例子中集成了mrcnn的代码，不过你还需要下载模型权重文件来进行物体识别，可以下载mask_rcnn_coco.h5，它可以识别80种物体。

安装依赖

Mask_RCNN源码和我的例子中都有一个requirements.txt文件，安装其中的依赖即可

pip install -r requirement.txt

如何使用Mask R-CNN

创建Config类

import mrcnn.config

class SampleConfig(mrcnn.config.Config):
    # Give the configuration a recognizable name
    NAME = "SampleRec"

    IMAGES_PER_GPU = 1
    GPU_COUNT = 1

    NUM_CLASSES = 80 + 1

Config类主要用来指定训练和预测时的配置，这里指定了单GPU处理图像数，GPU个数和最大可识别的种类数目

模型创建

model = mrcnn.model.MaskRCNN(mode = "inference", config=SampleConfig(), model_dir=os.getcwd())
model.load_weights("./models/mask_rcnn_coco.h5", by_name=True)

这里有2步，第一行初始化mrcnn的神经网络数据结构，第二行加载我们下载的权重文件

处理图片

首先使用opencv读取图片

img = cv2.imread(input_img)

然后进行处理

r = model.detect(images=[img], verbose=0)

处理返回值

mrcnn接受一个图片数组，返回的也是一个数组，取第一个

r = r[0]

接下来我们会用到r中的masks和rois参数，masks包含了每个检测目标的掩码，rois包含每个检测目标的矩形区域，利用这两个属性结合GrabCut进行抠图

masks = r['masks']
rois = r['rois']
obj_count = masks.shape[2]

for i in range(0, obj_count):
    print("process: " + str(i))
    roi_mask = masks[0:, 0:, i:i+1]
    roi_rect = rois[i]
    obj_img, mask, roi = cut_image(img, roi_mask, roi_rect)

cut_image会单独处理每个区域的抠图，实现如下

def cut_image(origin_img, mask, roi):
    roi_img = origin_img[roi[0]: roi[2], roi[1]: roi[3], :]
    roi_mask = mask[roi[0]: roi[2], roi[1]: roi[3], :]

    initial_mask = roi_mask.astype(np.uint8) + 2
    
    fg_model = np.zeros((1, 65), np.float64)
    bg_model = np.zeros((1, 65), np.float64)
    cv2.grabCut(roi_img, initial_mask, (0, 0, initial_mask.shape[0], initial_mask.shape[1]), bg_model, fg_model, 12, cv2.GC_INIT_WITH_MASK)

    final_img = cv2.cvtColor((initial_mask & 1) * roi_img, cv2.COLOR_RGB2BGR)
    return final_img, initial_mask, roi

首先通过roi扣出图片和掩码相应矩形区域的图，然后通过掩码来进行grabCut。掩码默认是布尔值，我把它转换成了uint8并加2，这样可以让掩码的像素值达到grabCut的标准，2表示疑似背景，3表示疑似前景。

Mask R-CNN + GrabCut 自动抠图

概述