【Pytorch框架实战】之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]

二、代码

import torch
import cv2 as cv
import torchvision.transforms as transforms
import torchvision
import numpy as np

# 调用模型
model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
model.eval()

# 预处理
preprocess = transforms.Compose([transforms.ToTensor()])

# 使用GPU
train_on_gpu = torch.cuda.is_available()

if train_on_gpu:
    model.cuda()

COCO_INSTANCE_CATEGORY_NAMES = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
    'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
    'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]


if __name__ == "__main__":
    frame = cv.imread("demo_img13.jpg")
    blob = preprocess(frame)
    c, h, w = blob.shape
    input_x = blob.view(1, c, h, w)
    output = model(input_x.cuda())[0]
    boxes = output['boxes'].cpu().detach().numpy()
    scores = output['scores'].cpu().detach().numpy()
    labels = output['labels'].cpu().detach().numpy()
    masks = output['masks'].cpu().detach().numpy()
    index = 0
    color_mask = np.zeros((h, w, c), dtype=np.uint8)
    mv = cv.split(color_mask)
    # 循环遍历
    for x1, y1, x2, y2 in boxes:
        if scores[index] > 0.5:
            cv.rectangle(frame, (np.int32(x1), np.int32(y1)),
                         (np.int32(x2), np.int32(y2)), (0, 255, 255), 1, 8, 0)
            mask = np.squeeze(masks[index] > 0.5)
            np.random.randint(0, 256)
            mv[2][mask == 1], mv[1][mask == 1], mv[0][mask == 1] = \
                [np.random.randint(0, 256), np.random.randint(0, 256), np.random.randint(0, 256)]
            label_id = labels[index]
            label_txt = COCO_INSTANCE_CATEGORY_NAMES[label_id]
            cv.putText(frame, label_txt, (np.int32(x1), np.int32(y1)), cv.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 1)
        index += 1
    color_mask = cv.merge(mv)
    result = cv.addWeighted(frame, 0.5, color_mask, 0.5, 0)
    cv.imwrite("demo_img13_test.png", result)

三、结果