輕松學(xué)Pytorch –Mask-RCNN圖像實(shí)例分割

Faster-RCNN網(wǎng)絡(luò)主要由三個部分組成分別是backbone的卷積網(wǎng)絡(luò)、實(shí)現(xiàn)Boxes選擇的區(qū)域推薦網(wǎng)絡(luò)RPN、最終的分類回歸。Mask-RCNN簡單說就是在RPN之后得到對齊ROI對齊區(qū)域，完成了一個全卷積的像素分割分支，Mask-RCNN的網(wǎng)絡(luò)結(jié)構(gòu)如下：

在推理階段，模型輸出下列字典選項：

boxes：預(yù)測矩形的左上角與右下角坐標(biāo)(x1,y1,x2,y2) [Nx4]
labels: 預(yù)測每個對象標(biāo)簽
scores：預(yù)測每個對象的得分，在0~1之間，大于閾值T的即為預(yù)測輸出
masks：預(yù)測每個實(shí)例對象的mask，mask>0.5作為最終分類mask。[Nx1xHxW]

Pytorch中使用Mask-RCNN實(shí)現(xiàn)實(shí)例分割，是基于torchvision的預(yù)訓(xùn)練模型庫，首先需要下載預(yù)訓(xùn)練模型，并檢查是否可以支持GPU推理，相關(guān)的代碼如下：

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()

基于輸入圖像，實(shí)現(xiàn)Mask-RCNN模型推理預(yù)測，得到對象與實(shí)例分割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()

對推理預(yù)測得到四個輸出結(jié)果，分別進(jìn)行解析，其中score閾值為0.5，mask采用soft版本，對大于0.5分割為當(dāng)前對象像素，這部分的代碼實(shí)現(xiàn)如下：

 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

其中對實(shí)例分割的對象像素進(jìn)行隨機(jī)顏色填充，完成彩色mask圖像生成。最終把彩色mask圖像與輸入圖像進(jìn)行疊加，得到實(shí)例分割輸出結(jié)果如下：

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)

這里，我測試了三張圖像，結(jié)果分別如下：

這個是很久以前我寫過一個無人機(jī)的Mask-RCNN檢測時候，別人問我的問題，其實(shí)這個就是很簡單的OpenCV操作就可以很好的提取出來這些ROI圖像，代碼實(shí)現(xiàn)如下：

 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()

運(yùn)行結(jié)果如下：

其中有一些基本的OpenCV函數(shù)使用與操作，如果你對這些OpenCV基礎(chǔ)知識運(yùn)用想學(xué)習(xí)一波，直接免費(fèi)看我B站 OpenCV4 30講，點(diǎn)擊這里：

好消息！
小白學(xué)視覺知識星球
開始面向外開放啦??????



下載1：OpenCV-Contrib擴(kuò)展模塊中文版教程
在「小白學(xué)視覺」公眾號后臺回復(fù)：擴(kuò)展模塊中文教程，即可下載全網(wǎng)第一份OpenCV擴(kuò)展模塊教程中文版，涵蓋擴(kuò)展模塊安裝、SFM算法、立體視覺、目標(biāo)跟蹤、生物視覺、超分辨率處理等二十多章內(nèi)容。

下載2：Python視覺實(shí)戰(zhàn)項目52講
在「小白學(xué)視覺」公眾號后臺回復(fù)：Python視覺實(shí)戰(zhàn)項目，即可下載包括圖像分割、口罩檢測、車道線檢測、車輛計數(shù)、添加眼線、車牌識別、字符識別、情緒檢測、文本內(nèi)容提取、面部識別等31個視覺實(shí)戰(zhàn)項目，助力快速學(xué)校計算機(jī)視覺。

下載3：OpenCV實(shí)戰(zhàn)項目20講
在「小白學(xué)視覺」公眾號后臺回復(fù)：OpenCV實(shí)戰(zhàn)項目20講，即可下載含有20個基于OpenCV實(shí)現(xiàn)20個實(shí)戰(zhàn)項目，實(shí)現(xiàn)OpenCV學(xué)習(xí)進(jìn)階。

交流群

歡迎加入公眾號讀者群一起和同行交流，目前有SLAM、三維視覺、傳感器、自動駕駛、計算攝影、檢測、分割、識別、醫(yī)學(xué)影像、GAN、算法競賽等微信群（以后會逐漸細(xì)分），請掃描下面微信號加群，備注：”昵稱+學(xué)校/公司+研究方向“，例如：”張三 + 上海交大 + 視覺SLAM“。請按照格式備注，否則不予通過。添加成功后會根據(jù)研究方向邀請進(jìn)入相關(guān)微信群。請勿在群內(nèi)發(fā)送廣告，否則會請出群，謝謝理解~