首页 文章详情

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

小白学视觉 | 25 2022-11-24 21:11 0 0 0
UniSMS (合一短信)

点击上方小白学视觉”,选择加"星标"或“置顶

重磅干货,第一时间送达



前面介绍了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()])
# 使用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)), (0255255), 180)
8        mask = np.squeeze(masks[index] > 0.5)
9        np.random.randint(0256)
10        mv[2][mask == 1], mv[1][mask == 1], mv[0][mask == 1] = \
11            [np.random.randint(0256), np.random.randint(0256), np.random.randint(0256)]
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, (00255), 1)
16    index += 1

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

1color_mask = cv.merge(mv)
2result = cv.addWeighted(frame, 0.5, color_mask, 0.50)
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[:,:,:] = (2550255)
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] = (000)
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()

运行结果如下:

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

好消息!

小白学视觉知识星球

开始面向外开放啦👇👇👇




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

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

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

交流群


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


good-icon 0
favorite-icon 0
收藏
回复数量: 0
    暂无评论~~
    Ctrl+Enter