使用 Mask R-CNN 的實例分割

使用 PyTorch 進行影像深度學習

Michal Oleszak

Machine Learning Engineer

Faster R-CNN

Faster R-CNN 架構示意圖

Mask R-CNN

Mask R-CNN 架構示意圖

PyTorch 的預訓練 Mask R-CNN

from torchvision.models.detection import \
maskrcnn_resnet50_fpn


model = maskrcnn_resnet50_fpn(pretrained=True)
model.eval()


image = Image.open("cat_and_laptop.jpg")
transform = transforms.Compose([
    transforms.ToTensor()
])
image_tensor = transform(image).unsqueeze(0)


with torch.no_grad():
    prediction = model(image_tensor)

匯入 Mask R-CNN 模型
載入預先訓練模型
載入測試影像並轉為張量

一隻貓坐在筆電旁的照片

將影像張量送入模型

模型輸出

標籤

prediction[0]["labels"]

tensor([
    17, 73, 76, 73, 67, 42, 63, 84,73, 65, 
    17, 73, 73, 73, 84, 72, 76, 76,17, 15
])

類別名稱

print(class_names[17], class_names[73])

cat laptop

類別機率

prediction[0]["scores"]

tensor([
    0.9981, 0.9672, 0.9061, 0.6893, 0.3729, 
    ..., 
    0.0745, 0.0705, 0.0623, 0.0610, 0.0508
])

遮罩（masks）

prediction[0]["masks"]

tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
          ...]]]])

Soft masks（軟遮罩）

遮罩的唯一值

prediction[0]["masks"].unique()

tensor([0.0000e+00, 5.9713e-08, ..., 
        9.9989e-01, 9.9990e-01])

Mask R-CNN 的遮罩：
- 介於 0 與 1 之間的數值
- 表示模型判斷每個像素屬於目標的信心
- 比二元遮罩提供更細緻的資訊
- 需要時可用閾值轉為二元遮罩

顯示軟遮罩

masks = prediction[0]["masks"]
labels = prediction[0]["labels"]


for i in range(2):
    plt.imshow(image)

    plt.imshow(
        masks[i, 0],
        cmap="jet",
        alpha=0.5,
    )

    plt.title(
        f"Object: {class_names[labels[i]]}"
    )
    plt.show()

從預測中取出遮罩與標籤
迭代前兩個物件並繪出原圖
對每個物件疊加半透明遮罩
加上標題並顯示

顯示軟遮罩

覆疊在影像上的實例分割遮罩

一起來練習吧！

使用 PyTorch 進行影像深度學習