pytorch下训练自己的Mask R-CNN实例分割模型

概述

最近想学习下Mark R-CNN模型做实例分割,想快速体验下,所以暂时不想花时间去看源码的具体实现细节.RCNN系列模型介绍推荐看这篇博客目标检测模型概述.里面主要介绍了目标检测模型的发展.后来看到一篇博客链接看到较新版本的torchvision里面有现成的Mask R-CNN模型可以直接使用,且程序里面注释有demo教你快速使用.因此自己就马上体验下啦(文尾有全部代码,utils.py,transforms.py,coco_eval.py,engine.py,coco_utils.py从这里下载链接).
本机环境:系统ubuntu16.4,GTX 960,cuda10.0,cudnn7.6,pytorch1.2.0,torchvision0.4.0

测试预训练模型

导入相应包

import os
import torch
import numpy as np
import torch.utils.data
from PIL import Image
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
import utils
import transforms as T
from engine import train_one_epoch, evaluate
import cv2
from random import shuffle

为了后面的检测结果显示,定义coco数据集的80类对象对应的显示颜色

class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
               'bus', 'train', 'truck', 'boat', 'traffic light',
               'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
               'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
               'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
               'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
               'kite', 'baseball bat', 'baseball glove', 'skateboard',
               'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
               'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
               'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
               'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
               'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
               'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
               'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
               'teddy bear', 'hair drier', 'toothbrush']
colors=[[np.random.randint(0,255),np.random.randint(0,255),np.random.randint(0,255)]for i in range(100)]
#为了最终实例分割显示明显,定义常见类别为深色
colors[1]=[255,0,0] #person
colors[2]=[0,255,0] #bicycle
colors[3]=[0,0,255] #car
colors[4]=[255,255,0]#motorcycle

直接用model=torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)载入预训练模型,这个函数会下载模型文件的,网速嘛…自己FQ吧.模型输出为一个列表,每个元素是每张图的检测结果–dict,由masks,labels,scores三个key得到最终结果.demo程序如下(分数阈值取的0.5).

def demo():
    img_dir='/home/zyk/dataset/PennFudanPed/PNGImages'
    #img_dir='/home/zyk/datasets/PennFudanPed/PNGImages'
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    # load an instance segmentation model pre-trained on COCO
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)

    model.to(device)
    imgs=os.listdir(img_dir)
    shuffle(imgs)
    for i in range(50):
        imgsrc=cv2.imread(os.path.join(img_dir,imgs[i]))
        all_cls_mask_color = np.zeros_like(imgsrc)
        all_cls_mask_index=np.zeros_like(imgsrc)
        img = imgsrc / 255.
        img=np.transpose(img, (2, 0, 1))
        img=torch.tensor(img,dtype=torch.float)
        # put the model in evaluation mode
        model.eval()

        with torch.no_grad():
            prediction = model([img.to(device)])
            #print(prediction)#可以简单传入一个tensor,看看输出返回结果
            scores =prediction[0]['scores']
            for idx,score in enumerate(scores):
                if score>0.5:
                    mask=prediction[0]['masks'][idx][0].cpu().numpy()
                    mask=mask>0.5
                    cls_id=prediction[0]['labels'][idx].item()
                    all_cls_mask_color[mask]=colors[cls_id]
                    all_cls_mask_index