本文介紹 Detectron2訓練自己的實例分割數據集
Detectron2訓練自己的實例分割數據集
This article was original written by XRBLS, welcome re-post, first come with https://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？
好在有這個需求的人不少，這值得我們開發一個工具來轉它。我們維護的現成的工具腳本可以在這裏找到：
https://github.com/jinfagang/cityscapestococo
實際上我們將數據轉到coco之後，發現效果我們想象的還好，因爲這些標註都非常精準，並且包含了我們感興趣的常用類別，比如car，truck，rider，cyclist，motor-cycle等。這些步驟與我們上面訓練堅果數據集的步驟差不多，大家仿造步驟來，首先可視化數據確保我們的數據是準確無誤的。那麼接下里的事情就很簡單了。
看起來似乎還不錯，這些類別衆多而精準，接下來開始train。
這個instance segmention模型訓練完成之後是一種什麼樣的效果呢？模型還在訓練，我們等下一期技術文章發佈之後，給大家分享一下結果，同時，告訴大家一個好消息：
- 類似於MaskRCNN這樣的模型速度是很難在實際場景中用起來的，原因也很簡單，我們都需要realtime的東西，這速度差強人意，但是否有人嘗試更輕便的backbone呢？不得而知，不過好消息是實際上通過一些加速框架，是可以把MaskRCNN加速到realtime的。而這些，我們等到detectron2go發佈之後，或許會看得到一些好消息。
本文首發於MANA AI社區，如果你對AI感興趣，歡迎加入我們的社區一起交流：http://t.manaai.cn