Detectron2訓練本身的實例分割數據集

本文介紹 Detectron2訓練本身的實例分割數據集

Detectron2訓練本身的實例分割數據集

This article was original written by Jin Tian, welcome re-post, first come with jinfagang.github.io . but please keep this copyright info, thanks, any question could be asked via wechat: jintianiloveu

本文介紹如何構造本身的類coco數據集，並用detectron2來訓練並預測。實際上detectron2出來已經有幾天了。但這個框架我的感受離真正工業使用還有點距離，首先第一點是很差部署，其次是相關的其餘模型導出支持比較差，好比不支持onnx，同時即使是導出的onnx模型也很難用上一些加速框架進行加速，好比不支持TensorRT。但若是你不是追求速度，用python作推理也是能夠的，而且最關鍵的是，你能夠用你的數據集訓練你的模型，或者是在這個框架上增長一些本身的修改，看看效果是否是更好之類。

首先看看這個如何來訓練本身的數據集。咱們今天要用的數據是這個：

wget https://github.com/Tony607/detectron2_instance_segmentation_demo/releases/download/V0.1/data.zip
複製代碼

這篇文章很大借鑑於原做者的一些嘗試，感興趣的朋友能夠給他github一個star，人家也不容易。這個data是一個很是很是適合來測試分割的一個微型數據集，小到什麼程度？只有那麼25張圖片。。

類別大概是:

cats: [{'supercategory': 'date', 'id': 1, 'name': 'date'}, {'supercategory': 'fig', 'id': 2, 'name': 'fig'}, {'supercategory': 'hazelnut', 'id': 3, 'name': 'hazelnut'}]
複製代碼

（這裏date的意思是棗椰子，fig的意思是無花果，hazelnut是榛子。)

若是你下載好了數據，那麼基本上咱們能夠開始了。你們能夠看到這個數據集還有一個trainval.json，徹底是按照coco的標註方式來標註的。coco的json格式也是目前比較通用的是instance segmentation 或者是Panoptic segmentation標註格式。

Setup Detectron2

關於如何安裝detectron2這裏不展開敘述，你們能夠按照github給予的步驟來。這裏給你們幾點提醒：

• 要python3.6;
• 要pytorch1.3.

其餘的沒有了。

訓練

首先看一下這個很是小很是迷你的數據集，在100次訓練以後的效果：

能夠看出來，這個效果仍是很是不錯的啊。訓練的腳本也很是簡單：

import random
from detectron2.utils.visualizer import Visualizer
from detectron2.data.catalog import MetadataCatalog, DatasetCatalog
import fruitsnuts_data
import cv2
from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
import os
from detectron2.engine.defaults import DefaultPredictor
from detectron2.utils.visualizer import ColorMode


fruits_nuts_metadata = MetadataCatalog.get("fruits_nuts")


if __name__ == "__main__":
    cfg = get_cfg()
    cfg.merge_from_file(
        "../../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
    )
    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
    print('loading from: {}'.format(cfg.MODEL.WEIGHTS))
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5   # set the testing threshold for this model
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
    cfg.DATASETS.TEST = ("fruits_nuts", )
    predictor = DefaultPredictor(cfg)

    data_f = './data/images/2.jpg'
    im = cv2.imread(data_f)
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                   metadata=fruits_nuts_metadata,
                   scale=0.8,
                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels
                   )
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    img = v.get_image()[:, :, ::-1]
    cv2.imshow('rr', img)
    cv2.waitKey(0)

複製代碼

你們可能以爲這個instance segmention訓練的太容易，那麼來挑戰難度大一點的？

確實，這個堅果數據集實在是太簡單了，簡單到咱們甚至打印不出任何訓練過程，一眨眼模型就訓練好了。那就來個難度大一點的吧，咱們將用Detectron2訓練Cityscapes的實例分割！

Detectron2 訓練Cityscapes實例分割

儘管在Detectron2裏面有Cityscapes的數據dataloader，但咱們並不打算用它，相反，咱們以爲使用統一化的coco數據集格式可以更好的讓數據fit。當咱們決定作這件事情的時候犯難了。如何從Cityscapes轉到coco？

好在有這個需求的人很多，這值得咱們開發一個工具來轉它。咱們維護的現成的工具腳本能夠在這裏找到：

github.com/jinfagang/c…

實際上咱們將數據轉到coco以後，發現效果咱們想象的還好，由於這些標註都很是精準，而且包含了咱們感興趣的經常使用類別，好比car，truck，rider，cyclist，motor-cycle等：

這是可視化的效果。這些步驟與咱們上面訓練堅果數據集的步驟差很少，你們仿造步驟來，首先可視化數據確保咱們的數據是準確無誤的。那麼接下里的事情就很簡單了。

看起來彷佛還不錯，這些類別衆多而精準，接下來開始train：

看起來彷佛還不錯。如今訓練一個instance segmentation 簡直是太簡單了。

那麼這個instance segmention模型訓練完成以後是一種什麼樣的效果呢？模型還在訓練，咱們等下一期技術文章發佈以後，給你們分享一下結果，同時，告訴你們一個好消息：

1. 相似於MaskRCNN這樣的模型速度是很難在實際場景中用起來的，緣由也很簡單，咱們都須要realtime的東西，這速度差強人意，可是否有人嘗試更輕便的backbone呢？不得而知，不過好消息是實際上經過一些加速框架，是能夠把MaskRCNN加速到realtime的。而這些，咱們等到detectron2go發佈以後，或許會看獲得一些好消息。

本文首發於MANA AI社區，若是你對AI感興趣，歡迎加入咱們的社區一塊兒交流：t.manaai.cn