轻松学Pytorch –Mask-RCNN图像实例分割
轻松学Pytorch –Mask-RCNN图像实例分割
前面介绍了torchvison框架下Faster-RCNN对象检测模型使用与自定义对象检测的数据集制作与训练。在计算机视觉所要面对的任务中，最常见的就是对象检测、图像语义分割跟实例分割，torchvision支持Mask-RCNN模型的调用与自定义数据训练，可以同时实现对象检测与实例分割任务。本文主要跟大家分享一下如何使用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()])
# 使用GPU
train_on_gpu = torch.cuda.is_available()
if train_on_gpu:
model.cuda()基于输入图像，实现Mask-RCNN模型推理预测，得到对象与实例分割mask的代码如下：
1 frame = cv.imread("D:/images/master.jpg")
2 blob = transform(frame)
3 c, h, w = blob.shape
4 input_x = blob.view(1, c, h, w)
5 output = model(input_x.cuda())[0]
6 boxes = output['boxes'].cpu().detach().numpy()
7 scores = output['scores'].cpu().detach().numpy()
8 labels = output['labels'].cpu().detach().numpy()
9 masks = output['masks'].cpu().detach().numpy()对推理预测得到四个输出结果，分别进行解析，其中score阈值为0.5，mask采用soft版本，对大于0.5分割为当前对象像素，这部分的代码实现如下：
1 index = 0
2 color_mask = np.zeros((h, w, c), dtype=np.uint8)
3 mv = cv.split(color_mask)
4 for 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图像与输入图像进行叠加，得到实例分割输出结果如下：
1 color_mask = cv.merge(mv)
2 result = cv.addWeighted(frame, 0.5, color_mask, 0.5, 0)
3 cv.imshow("intances segmentation demo", result)
4 cv.imwrite("D:/master_test.png", result)这里，我测试了三张图像，结果分别如下：
Mask-RCNN实例分割对象提取与背景替换
这个是很久以前我写过一个无人机的Mask-RCNN检测时候，别人问我的问题，其实这个就是很简单的OpenCV操作就可以很好的提取出来这些ROI图像，代码实现如下：
1 # 简单背景替换
2 back_ground = np.zeros_like(frame)
3 back_ground[:,:,:] = (255, 0, 255)
4 for 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)
9 temp = cv.add(back_ground, frame, mask=mv[0])
10 dst = cv.add(back_ground, temp)
11 cv.imshow("background replacement", dst)
12 cv.waitKey(0)
13 cv.destroyAllWindows()运行结果如下：
