輕鬆學Pytorch –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()])
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)

這裡，我測試了三張圖像，結果分別如下：
這個是很久以前我寫過一個無人機的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()

運行結果如下：

其中有一些基本的OpenCV函數使用與操作，如果你對這些OpenCV基礎知識運用想學習一波，直接免費看我B站 OpenCV4 30講，點擊這裡：
全程實操 | 最新版OpenCV4.4免費視頻課程送給大家

 推薦閱讀 
輕鬆學Pytorch–環境搭建與基本語法
Pytorch輕鬆學-構建淺層神經網絡
輕鬆學pytorch-構建卷積神經網絡
輕鬆學Pytorch –構建循環神經網絡
輕鬆學Pytorch-使用卷積神經網絡實現圖像分類
輕鬆學Pytorch-自定義數據集製作與使用
輕鬆學Pytorch-Pytorch可視化
輕鬆學Pytorch–Visdom可視化
輕鬆學Pytorch – 全局池化層詳解
輕鬆學Pytorch – 人臉五點landmark提取網絡訓練與使用
輕鬆學Pytorch – 年齡與性別預測
輕鬆學Pytorch –車輛類型與顏色識別
輕鬆學Pytorch-全卷積神經網絡實現表情識別
使用OpenVINO加速Pytorch表情識別模型
輕鬆學pytorch – 使用多標籤損失函數訓練卷積網絡
輕鬆學Pytorch-使用ResNet50實現圖像分類
OpenCV4.4 + YOLOv4 真的可以運行了…..
輕鬆學Pytorch –使用torchvision實現對象檢測
輕鬆學Pytorch-實現自定義對象檢測器
從Pytorch 的ONNX到OpenVINO中IR中間層

伏久者，飛必高
開先者，謝獨早