在Pytorch中使用Mask R-CNN進行實例分割操作

在這篇文章中，我們將討論mask R-CNN背後的一些理論，以及如何在PyTorch中使用預訓練的mask R-CNN模型。

1.語義分割、目標檢測和實例分割

之前已經介紹過：

1、語義分割：在語義分割中，我們分配一個類標籤（例如。狗、貓、人、背景等）對圖像中的每個像素。

2、目標檢測：在目標檢測中，我們將類標籤分配給包含對象的包圍框。

一個非常自然的想法是把兩者結合起來。我們只想在一個對象周圍識別一個包圍框，並且找到包圍框中的哪些像素屬於對象。 換句話說，我們想要一個掩碼，它指示（使用顏色或灰度值）哪些像素屬於同一對象。 產生上述掩碼的一類算法稱為實例分割算法。mask R-CNN就是這樣一種算法。

實例分割和語義分割有兩種不同

1、在語義分割中，每個像素都被分配一個類標籤，而在實例分割中則不是這樣。

2、在語義分割中，我們不區分同一類的實例。例如，語義分割中屬於“Person”類的所有像素都將在掩碼中分配相同的顏色/值。在實例分割中，它們被分配到不同的值，我們能夠告訴它們哪個像素對應於哪個人。 要了解更多關於圖像分割的信息，請查看我們已經詳細解釋過的帖子。

Mask R-CNN結構

mask R-CNN的網絡結構是我們以前討論過的FasterR-CNN的擴展。

回想一下，faster R-CNN架構有以下組件

卷積層：輸入圖像經過幾個卷積層來創建特徵圖。如果你是初學者，把卷積層看作一個黑匣子，它接收一個3通道的輸入圖像，並輸出一個空間維數小得多（7×7），但通道非常多（512）的“圖像”。

區域提案網絡（RPN）。卷積層的輸出用於訓練一個網絡，該網絡提取包圍對象的區域。

分類器：同樣的特徵圖也被用來訓練一個分類器，該分類器將標籤分配給框內的對象。

此外，回想一下，FasterR-CNN 比 Fast R-CNN更快，因為特徵圖被計算一次，並被RPN和分類器重用。 mask R-CNN將這個想法向前推進了一步。除了向RPN和分類器提供特徵圖外，mask R-CNN還使用它來預測邊界框內對象的二值掩碼。 研究 MaskR-CNN的掩碼預測部分的一種方法是，它是一個用於語義分割的完全卷積網絡（FCN）。唯一的區別是在mask R-CNN裡，FCN被應用於邊界框，而且它與RPN和分類器共享卷積層。 下圖顯示了一個非常高層次的架構。

2.在PyTorch中使用mask R-CNN[代碼]

在本節中，我們將學習如何在PyTorch中使用預先訓練的MaskR-CNN模型。

2.1.輸入和輸出

mask R-CNN模型期望的輸入是張量列表，每個張量的類型為（n，c，h，w），元素在0-1範圍內。圖像的大小隨意。

n是圖像的個數

c為通道數 RGB圖像為3

h是圖像的高度

w是圖像的寬度

模型返回 ：

包圍框的座標

模型預測的存在於輸入圖像中的類的標籤以及對應標籤的分數

標籤中每個類的掩碼。

2.2 預訓練模型

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
 model.eval()

2.3 模型的預測  

  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'  ]     def get_prediction(img_path, threshold):   img = Image.open(img_path)   transform = T.Compose([T.ToTensor()])   img = transform(img)   pred = model([img])   print('pred')   print(pred)   pred_score = list(pred[0]['scores'].detach().numpy())   pred_t = [pred_score.index(x) for x in pred_score if x>threshold][-1]   print("masks>0.5")   print(pred[0]['masks']>0.5)   masks = (pred[0]['masks']>0.5).squeeze().detach().cpu().numpy()   print("this is masks")   print(masks)   pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())]   pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]   masks = masks[:pred_t+1]   pred_boxes = pred_boxes[:pred_t+1]   pred_class = pred_class[:pred_t+1]   return masks, pred_boxes, pred_class

代碼功能如下：

從圖像路徑中獲取圖像

使用PyTorch變換將圖像轉換為圖像張量

通過模型傳遞圖像以得到預測結果

從模型中獲得掩碼、預測類和包圍框座標

每個預測對象的掩碼從一組11個預定義的顏色中隨機給出顏色，以便在輸入圖像上將掩碼可視化。

  def random_colour_masks(image):   colours = [[0, 255, 0],[0, 0, 255],[255, 0, 0],[0, 255, 255],[255, 255, 0],[255, 0, 255],[80, 70, 180],[250, 80, 190],[245, 145, 50],[70, 150, 250],[50, 190, 190]]   r = np.zeros_like(image).astype(np.uint8)   g = np.zeros_like(image).astype(np.uint8)   b = np.zeros_like(image).astype(np.uint8)   r[image == 1], g[image == 1], b[image == 1] = colours[random.randrange(0,10)]   coloured_mask = np.stack([r, g, b], axis=2)   return coloured_mask

代碼中有一些打印信息幫助分析處理過程

2.4 實例分割工作流程

  def instance_segmentation_api(img_path, threshold=0.5, rect_th=3, text_size=3, text_th=3):   masks, boxes, pred_cls = get_prediction(img_path, threshold)   img = cv2.imread(img_path)   img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)   for i in range(len(masks)):    rgb_mask = random_colour_masks(masks[i])    img = cv2.addWeighted(img, 1, rgb_mask, 0.5, 0)    cv2.rectangle(img, boxes[i][0], boxes[i][1],color=(0, 255, 0), thickness=rect_th)    cv2.putText(img,pred_cls[i], boxes[i][0], cv2.FONT_HERSHEY_SIMPLEX, text_size, (0,255,0),thickness=text_th)   plt.figure(figsize=(20,30))   plt.imshow(img)   plt.xticks([])   plt.yticks([])   plt.show()

掩碼、預測類和邊界框是通過get_prediction獲得的。

每個掩碼從11種顏色中隨機給出顏色。 每個掩碼按比例1：0.5被添加到圖像中，使用了opencv。

包圍框是用cv2.rectangle繪製的，上面有類名。

顯示最終輸出

完整代碼如下：

  from PIL import Image  import matplotlib.pyplot as plt  import torch  import torchvision.transforms as T  import torchvision  import torch  import numpy as np  import cv2  import random     model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)  model.eval()  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'  ]  def get_prediction(img_path, threshold):   img = Image.open(img_path)   transform = T.Compose([T.ToTensor()])   img = transform(img)   pred = model([img])   print('pred')   print(pred)   pred_score = list(pred[0]['scores'].detach().numpy())   pred_t = [pred_score.index(x) for x in pred_score if x>threshold][-1]   print("masks>0.5")   print(pred[0]['masks']>0.5)   masks = (pred[0]['masks']>0.5).squeeze().detach().cpu().numpy()   print("this is masks")   print(masks)   pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())]   pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]   masks = masks[:pred_t+1]   pred_boxes = pred_boxes[:pred_t+1]   pred_class = pred_class[:pred_t+1]   return masks, pred_boxes, pred_class     def random_colour_masks(image):   colours = [[0, 255, 0],[0, 0, 255],[255, 0, 0],[0, 255, 255],[255, 255, 0],[255, 0, 255],[80, 70, 180],[250, 80, 190],[245, 145, 50],[70, 150, 250],[50, 190, 190]]   r = np.zeros_like(image).astype(np.uint8)   g = np.zeros_like(image).astype(np.uint8)   b = np.zeros_like(image).astype(np.uint8)   r[image == 1], g[image == 1], b[image == 1] = colours[random.randrange(0,10)]   coloured_mask = np.stack([r, g, b], axis=2)   return coloured_mask     def instance_segmentation_api(img_path, threshold=0.5, rect_th=3, text_size=3, text_th=3):   masks, boxes, pred_cls = get_prediction(img_path, threshold)   img = cv2.imread(img_path)   img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)   for i in range(len(masks)):    rgb_mask = random_colour_masks(masks[i])    img = cv2.addWeighted(img, 1, rgb_mask, 0.5, 0)    cv2.rectangle(img, boxes[i][0], boxes[i][1],color=(0, 255, 0), thickness=rect_th)    cv2.putText(img,pred_cls[i], boxes[i][0], cv2.FONT_HERSHEY_SIMPLEX, text_size, (0,255,0),thickness=text_th)   plt.figure(figsize=(20,30))   plt.imshow(img)   plt.xticks([])   plt.yticks([])   plt.show()

2.5 示例

示例1：以小雞為例，會識別為鳥類

instance_segmentation_api('chicken.jpg')

輸入圖像：

輸出結果：

處理過程中的打印信息：

  pred  [{'boxes': tensor([[176.8106, 125.6315, 326.8023, 400.4467],      [427.9514, 130.5811, 584.2725, 403.1004],      [289.9471, 169.1313, 448.9896, 410.0000],      [208.7829, 140.7450, 421.3497, 409.0258],      [417.7833, 137.5480, 603.2806, 405.6804],      [174.3626, 132.7247, 330.4560, 404.6956],      [291.6709, 165.4233, 447.1820, 401.7686],      [171.9978, 114.4133, 336.9987, 410.0000],      [427.0312, 129.5812, 584.2130, 405.4166]], grad_fn=), 'labels': tensor([16, 16, 16, 16, 20, 20, 20, 18, 18]), 'scores': tensor([0.9912, 0.9910, 0.9894, 0.2994, 0.2108, 0.1995, 0.1795, 0.1655, 0.0516],      grad_fn=), 'masks': tensor([[[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]],         [[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]],         [[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]],         ...,         [[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]],         [[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]],         [[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]]], grad_fn=)}]  masks>0.5  tensor([[[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]],         [[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]],         [[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]],         ...,         [[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]],         [[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]],         [[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]]])  this is masks  [[[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]      [[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]      [[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]      ...      [[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]      [[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]      [[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]]

masks = (pred[0]['masks']>0.5).squeeze().detach().cpu().numpy()使masks變為[n x h x w],且元素為bool值，為後續指定隨機顏色做了準備，r[image == 1], g[image == 1], b[image == 1] = colours[random.randrange(0,10)]，將掩碼列表中屬於實際對象的區域變成隨機彩色，其餘部分仍為0.這些代碼充分展示了python中高級切片的魔力，當然用到的是numpy和torch.tensor裡的功能。

示例2：棕熊

instance_segmentation_api('bear.jpg', threshold=0.8)

輸入圖像：

輸出圖像：

打印信息：

  pred  [{'boxes': tensor([[ 660.3120, 340.5351, 1235.1614, 846.9672],      [ 171.7622, 426.9127, 756.6520, 784.9360],      [ 317.9777, 184.6863, 648.0856, 473.6469],      [ 283.0787, 200.8575, 703.7324, 664.4083],      [ 354.9362, 308.0444, 919.0403, 812.0120]], grad_fn=), 'labels': tensor([23, 23, 23, 23, 23]), 'scores': tensor([0.9994, 0.9994, 0.9981, 0.5138, 0.0819], grad_fn=), 'masks': tensor([[[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]],         [[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]],         [[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]],         [[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]],         [[[0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       ...,       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.],       [0., 0., 0., ..., 0., 0., 0.]]]], grad_fn=)}]  masks>0.5  tensor([[[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]],         [[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]],         [[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]],         [[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]],         [[[False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       ...,       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False],       [False, False, False, ..., False, False, False]]]])  this is masks  [[[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]      [[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]      [[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]      [[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]      [[False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]   ...   [False False False ... False False False]   [False False False ... False False False]   [False False False ... False False False]]]

3、GPU與CPU時間對比

  def check_inference_time(image_path, gpu=False):    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)    model.eval()    img = Image.open(image_path)    transform = T.Compose([T.ToTensor()])    img = transform(img)    if gpu:      model.cuda()      img = img.cuda()    else:      model.cpu()      img = img.cpu()    start_time = time.time()    pred = model([img])    end_time = time.time()    return end_time-start_time     cpu_time = sum([check_inference_time('./people.jpg', gpu=False) for _ in range(5)])/5.0  gpu_time = sum([check_inference_time('./people.jpg', gpu=True) for _ in range(5)])/5.0  print('
Average Time take by the model with GPU = {}s
Average Time take by the model with CPU = {}s'.format(gpu_time, cpu_time))

結果：

Average Time take by the model with GPU = 0.5736178874969482s,
 Average Time take by the model with CPU = 10.966966199874879s