首先，匯入包。

import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt

#設定根目錄
ROOT_DIR = os.path.abspath("../")

#匯入Mask R-CNN
sys.path.append(ROOT_DIR)
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
#匯入coco的配置
sys.path.append(os.path.join(ROOT_DIR, "samples/coco/"))
import coco
 

%matplotlib inline 

#設定寫log的目錄以及儲存訓練model的目錄
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

#預訓練模型的目錄
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
#如果檔案不存在，則下載COCO預訓練權重
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)

#測試圖片目錄
IMAGE_DIR = os.path.join(ROOT_DIR, "images")
 

這裡將會使用一個在MS_COCO資料集上訓練好的模型。這個模型的配置資訊位於coco.py檔案的CocoConfig類中。

在進行預測時，需要對這個配置做一些小的改動以匹配當前的任務。具體做法是，繼承CocoConfig類並更改你需要的屬性。

class InferenceConfig(coco.CocoConfig):
    #將batch size設定為1，因為在預測時一次只需要處理一副影象
    #Batch size = GPU_COUNT * IMAGES_PER_GPU
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

config = InferenceConfig()
#將配置資訊打印出來
config.display()
 

接下來就建立模型物件並載入訓練好的權重。

#建立預測模型物件.
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

#載入MS-COCO訓練權重
model.load_weights(COCO_MODEL_PATH, by_name=True)

該模型分類物體並返回類別ID，類別ID以整型值來區分每一類物體。

有些資料集為類別賦予整型值以便於區分，但是有些資料集沒有這樣做。ID通常是順序連續的，但不絕對。為了增強一致性並且在同一時間可以訓練多個源資料，這裡的Dataset class為每個類分配了自己的整型ID。

所以要注意將class IDs對映到正確的class names。

這裡我們並不想下載所有的COCO資料集來測試這個demo，所以將所有的class names預先存在一個list中。class name的索引即代表了它們的ID。

# COCO Class names
#索引即代表ID.例如，要獲取teddy bear的ID
# use: class_names.index('teddy bear')
class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
               'bus', 'train', 'truck', 'boat', 'traffic light',
               'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
               'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
               'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
               'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
               'kite', 'baseball bat', 'baseball glove', 'skateboard',
               'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
               'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
               'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
               'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
               'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
               'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
               'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
               'teddy bear', 'hair drier', 'toothbrush']

一切準備好後，就可以開始預測。從測試資料夾內隨機選取一副影象。

#隨機選擇影象
file_names = next(os.walk(IMAGE_DIR))[2]
#讀取影象
image = skimage.io.imread(os.path.join(IMAGE_DIR, random.choice(file_names)))

#執行檢測
results = model.detect([image], verbose=1)

#視覺化檢測結果
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], 
                            class_names, r['scores'])

首先會打印出相關資訊。

Processing 1 images
image                    shape: (476, 640, 3)         min:    0.00000  max:  255.00000
molded_images            shape: (1, 1024, 1024, 3)    min: -123.70000  max:  120.30000
image_metas              shape: (1, 89)               min:    0.00000  max: 1024.00000

接著會顯示檢測結果，包含bounding box， probability和mask。