轻松学Pytorch –Mask-RCNN图像实例分割

来源 | OpenCV学堂

作者 | gloomyfish

本文主要跟大家分享一下如何使用mask-rcnn网络实现对象检测与实例分割，下一篇将会介绍如何制作数据集训练Mask-RCNN网络。

Mask-RCNN网络模型

Faster-RCNN网络主要由三个部分组成分别是backbone的卷积网络、实现Boxes选择的区域推荐网络RPN、最终的分类回归。Mask-RCNN简单说就是在RPN之后得到对齐ROI对齐区域，完成了一个全卷积的像素分割分支，Mask-RCNN的网络结构如下：

在推理阶段，模型输出下列字典选项：

boxes：预测矩形的左上角与右下角坐标(x1,y1,x2,y2) [Nx4]labels: 预测每个对象标签scores：预测每个对象的得分，在0~1之间，大于阈值T的即为预测输出masks：预测每个实例对象的mask，mask>0.5作为最终分类mask。[Nx1xHxW]

使用Mask-RCNN实现实例分割

Pytorch中使用Mask-RCNN实现实例分割，是基于torchvision的预训练模型库，首先需要下载预训练模型，并检查是否可以支持GPU推理，相关的代码如下：

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)model.eval()transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])# 使用GPUtrain_on_gpu = torch.cuda.is_available()if train_on_gpu:model.cuda()

基于输入图像，实现Mask-RCNN模型推理预测，得到对象与实例分割mask的代码如下：

1frame = cv.imread("D:/images/master.jpg")
2blob = transform(frame)
3c, h, w = blob.shape
4input_x = blob.view(1, c, h, w)
5output = model(input_x.cuda())[0]
6boxes = output['boxes'].cpu().detach().numpy()
7scores = output['scores'].cpu().detach().numpy()
8labels = output['labels'].cpu().detach().numpy()
9masks = output['masks'].cpu().detach().numpy()

对推理预测得到四个输出结果，分别进行解析，其中score阈值为0.5，mask采用soft版本，对大于0.5分割为当前对象像素，这部分的代码实现如下：

 1index = 0
 2color_mask = np.zeros((h, w, c), dtype=np.uint8)
 3mv = cv.split(color_mask)
 4for x1, y1, x2, y2 in boxes:
 5    if scores[index] > 0.5:
 6        cv.rectangle(frame, (np.int32(x1), np.int32(y1)),
 7                     (np.int32(x2), np.int32(y2)), (0, 255, 255), 1, 8, 0)
 8        mask = np.squeeze(masks[index] > 0.5)
 9        np.random.randint(0, 256)
10        mv[2][mask == 1], mv[1][mask == 1], mv[0][mask == 1] = \
11            [np.random.randint(0, 256), np.random.randint(0, 256), np.random.randint(0, 256)]
12
13        label_id = labels[index]
14        label_txt = coco_names[str(label_id)]
15        cv.putText(frame, label_txt, (np.int32(x1), np.int32(y1)), cv.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 1)
16    index += 1

其中对实例分割的对象像素进行随机颜色填充，完成彩色mask图像生成。最终把彩色mask图像与输入图像进行叠加，得到实例分割输出结果如下：

1color_mask = cv.merge(mv)
2result = cv.addWeighted(frame, 0.5, color_mask, 0.5, 0)
3cv.imshow("intances segmentation demo", result)
4cv.imwrite("D:/master_test.png", result)

这里，我测试了三张图像，结果分别如下：

Mask-RCNN实例分割对象提取与背景替换

这个是很久以前我写过一个无人机的Mask-RCNN检测时候，别人问我的问题，其实这个就是很简单的OpenCV操作就可以很好的提取出来这些ROI图像，代码实现如下：

 1# 简单背景替换
 2back_ground = np.zeros_like(frame)
 3back_ground[:,:,:] = (255, 0, 255)
 4for row in range(h):
 5    for col in range(w):
 6        b, g, r = color_mask[row, col]
 7        if b > 0 or g > 0 or r > 0:
 8            back_ground[row, col] = (0, 0, 0)
 9temp = cv.add(back_ground, frame, mask=mv[0])
10dst = cv.add(back_ground, temp)
11cv.imshow("background replacement", dst)
12cv.waitKey(0)
13cv.destroyAllWindows()

运行结果如下：