玩轉Facebook的maskrcnn-benchmark項目 2

maskrcnn-benchmark是Facebook開源的基準（benchmark）算法工程，其中包含檢測、分割和人體關鍵點等算法。

本系列包含兩篇：

• 第一篇 搭建環境；
• 第二篇 訓練和驗證；

---

訓練

使用maskrcnn-benchmark訓練模型，能夠參考。

數據集：

• 下載完整的COCO數據集：annotations、test201四、train201四、val2014；
• 下載FAIR提供的COCO小型驗證集：minival和valminusminival；

選擇訓練模板：e2e_mask_rcnn_R_50_FPN_1x.yaml，其中：

WEIGHT: "catalog://ImageNetPretrained/MSRA/R-50"  # 預訓練權重
DATASETS:  # 數據集
  TRAIN: ("coco_2014_train", "coco_2014_valminusminival")
  TEST: ("coco_2014_minival",)
MAX_ITER: 90000  # 最大訓練輪次
複製代碼

其餘參數的設置位置：maskrcnn_benchmark/config/defaults.py

如：

• _C.SOLVER.CHECKPOINT_PERIOD = 2500，保存輪次；
• _C.SOLVER.IMS_PER_BATCH = 16，訓練的batch_size；
• _C.OUTPUT_DIR = "./models"，模型輸出路徑；

指定GPU的數量：

export NGPUS=4
複製代碼

訓練模型：

python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_net.py --config-file "configs/e2e_mask_rcnn_R_50_FPN_1x.yaml"

nohup python -u -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_net.py --config-file "configs/e2e_mask_rcnn_R_50_FPN_1x.yaml" &
複製代碼

輸出的模型位於./models中，最後一個模型是model_0090000.pth。

---

測試

將配置文件e2e_mask_rcnn_R_50_FPN_1x.my.yaml的WEIGHT修改成模型路徑，如WEIGHT: "model_0090000.pth"。

測試邏輯以下：

def main():
    img_path = os.path.join(DATA_DIR, 'aoa-mina.jpeg')

    img = cv2.imread(img_path)
    print('[Info] img size: {}'.format(img.shape))
    
    config_file = "../configs/e2e_mask_rcnn_R_50_FPN_1x.my.yaml"

    cfg.merge_from_file(config_file)  # 設置配置文件
    cfg.merge_from_list(["MODEL.MASK_ON", True])
    cfg.merge_from_list(["MODEL.DEVICE", "cpu"])  # 指定爲CPU

    coco_demo = COCODemo(  # 建立模型文件
        cfg,
        # show_mask_heatmaps=True,
        min_image_size=800,
        confidence_threshold=0.7,
    )

    predictions = coco_demo.compute_prediction(img)
    top_predictions = coco_demo.select_top_predictions(predictions)
    show_mask(img, top_predictions, coco_demo)

    print('執行完成!')
複製代碼

show_mask顯示分割的效果，以matplotlib的圖像爲基礎，繪製多邊形框：

def show_mask(img, predictions, coco_demo):
    """ 顯示分割效果 :param img: numpy的圖像 :param predictions: 分割結果 :param coco_demo: 函數集 :return: 顯示分割效果 """
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    pylab.rcParams['figure.figsize'] = (8.0, 10.0)  # 圖片尺寸
    plt.imshow(img)  # 須要提早填充圖像
    # plt.show()

    extra_fields = predictions.extra_fields
    masks = extra_fields['mask']
    labels = extra_fields['labels']
    name_list = [coco_demo.CATEGORIES[l] for l in labels]

    seg_list = []
    for mask in masks:
        mask = torch.squeeze(mask)
        segmentation = binary_mask_to_polygon(mask, tolerance=2)[0]
        seg_list.append(segmentation)

    ax = plt.gca()
    ax.set_autoscale_on(False)

    polygons, color = [], []

    np.random.seed(37)

    for name, seg in zip(name_list, seg_list):
        c = (np.random.random((1, 3)) * 0.8 + 0.2).tolist()[0]

        poly = np.array(seg).reshape((int(len(seg) / 2), 2))
        c_x, c_y = get_center_of_polygon(poly)  # 計算多邊形的中心點

        # 0~26是大類別, 其他是小類別 同時 每一個標籤只繪製一次
        tc = c - np.array([0.5, 0.5, 0.5])  # 下降顏色
        tc = np.maximum(tc, 0.0)  # 最小值0

        plt.text(c_x, c_y, name, ha='left', wrap=True, color=tc,
                 bbox=dict(facecolor='white', alpha=0.5))  # 繪製標籤

        polygons.append(pylab.Polygon(poly))  # 繪製多邊形
        color.append(c)

    p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.4)  # 添加多邊形
    ax.add_collection(p)
    p = PatchCollection(polygons, facecolor='none', edgecolors=color, linewidths=2)  # 添加多邊形的框
    ax.add_collection(p)

    plt.axis('off')
    ax.get_xaxis().set_visible(False)  # this removes the ticks and numbers for x axis
    ax.get_yaxis().set_visible(False)  # this removes the ticks and numbers for y axis

    out_folder = os.path.join(ROOT_DIR, 'demo', 'out')
    mkdir_if_not_exist(out_folder)
    out_file = os.path.join(out_folder, 'test.png'.format())
    plt.savefig(out_file, bbox_inches='tight', pad_inches=0, dpi=200)

    plt.close()  # 避免全部圖像繪製在一塊兒
複製代碼

效果以下：

---

OK, that's all!