轻松学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的网络结构如下:

560d8c6ca7ad0e037dd4ea17a6e0cc17.png

在推理阶段,模型输出下列字典选项:

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的代码如下:

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)

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

01032d18f0c932816849a9b6a6e5d13e.png

81217209ab7e485c1f20ffecf689fbf2.png

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()

运行结果如下:

5732a97bebb7d933af1cbf7ac45c02a3.png

65d545f46c8dc9d42f675faa4c226b39.png

其中有一些基本的OpenCV函数使用与操作,如果你对这些OpenCV基础知识运用想学习一波,直接免费看我B站 OpenCV4 30讲,点击这里:

 
 

好消息!

小白学视觉知识星球

开始面向外开放啦👇👇👇

 
 

df9cda983183feb0a5bb1758da4c29bd.jpeg

下载1:OpenCV-Contrib扩展模块中文版教程

在「小白学视觉」公众号后台回复:扩展模块中文教程,即可下载全网第一份OpenCV扩展模块教程中文版,涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。


下载2:Python视觉实战项目52讲
在「小白学视觉」公众号后台回复:Python视觉实战项目,即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。


下载3:OpenCV实战项目20讲
在「小白学视觉」公众号后台回复:OpenCV实战项目20讲,即可下载含有20个基于OpenCV实现20个实战项目,实现OpenCV学习进阶。


交流群

欢迎加入公众号读者群一起和同行交流,目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群(以后会逐渐细分),请扫描下面微信号加群,备注:”昵称+学校/公司+研究方向“,例如:”张三 + 上海交大 + 视觉SLAM“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~
  • 2
    点赞
  • 12
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值