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