概述
本篇博客主要介绍如何使用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表示疑似前景。