佳云大脑|–Mask-RCNN图像实例分割,轻松学Pytorch

来源|OpenCV学堂
作者|gloomyfish
本文主要跟大家分享一下如何使用mask-rcnn网络实现对象检测与实例分割 , 下一篇将会介绍如何制作数据集训练Mask-RCNN网络 。
Mask-RCNN网络模型
Faster-RCNN网络主要由三个部分组成分别是backbone的卷积网络、实现Boxes选择的区域推荐网络RPN、最终的分类回归 。 Mask-RCNN简单说就是在RPN之后得到对齐ROI对齐区域 , 完成了一个全卷积的像素分割分支 , Mask-RCNN的网络结构如下:
佳云大脑|–Mask-RCNN图像实例分割,轻松学Pytorch
文章图片
在推理阶段 , 模型输出下列字典选项: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()iftrain_on_gpu:model.cuda()
基于输入图像 , 实现Mask-RCNN模型推理预测 , 得到对象与实例分割mask的代码如下:1frame=cv.imread("D:/images/master.jpg")2blob=transform(frame)3c,h,w=blob.shape4input_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=02color_mask=np.zeros((h,w,c),dtype=np.uint8)3mv=cv.split(color_mask)4forx1,y1,x2,y2inboxes:5ifscores[index]>0.5:6cv.rectangle(frame,(np.int32(x1),np.int32(y1)),7(np.int32(x2),np.int32(y2)),(0,255,255),1,8,0)8mask=np.squeeze(masks[index]>0.5)9np.random.randint(0,256)10mv[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)]1213label_id=labels[index]14label_txt=coco_names[str(label_id)]15cv.putText(frame,label_txt,(np.int32(x1),np.int32(y1)),cv.FONT_HERSHEY_PLAIN,1.0,(0,0,255),1)16index+=1
其中对实例分割的对象像素进行随机颜色填充 , 完成彩色mask图像生成 。 最终把彩色mask图像与输入图像进行叠加 , 得到实例分割输出结果如下:1color_mask=cv.merge(mv)2result=cv.addWeighted(frame,0.5,color_mask,0.5,0)3cv.imshow("intancessegmentationdemo",result)4cv.imwrite("D:/master_test.png",result)
这里 , 我测试了三张图像 , 结果分别如下:
佳云大脑|–Mask-RCNN图像实例分割,轻松学Pytorch
文章图片
佳云大脑|–Mask-RCNN图像实例分割,轻松学Pytorch
文章图片
Mask-RCNN实例分割对象提取与背景替换
这个是很久以前我写过一个无人机的Mask-RCNN检测时候 , 别人问我的问题 , 其实这个就是很简单的OpenCV操作就可以很好的提取出来这些ROI图像 , 代码实现如下:1#简单背景替换2back_ground=np.zeros_like(frame)3back_ground[:,:,:]=(255,0,255)4forrowinrange(h):5forcolinrange(w):6b,g,r=color_mask[row,col]7ifb>0org>0orr>0:8back_ground[row,col]=(0,0,0)9temp=cv.add(back_ground,frame,mask=mv[0])10dst=cv.add(back_ground,temp)11cv.imshow("backgroundreplacement",dst)12cv.waitKey(0)13cv.destroyAllWindows()


推荐阅读