tta部署_Detectron2 部署简单的训练项目

Detectron2 官方文档里的 Getting Started 提供了两种使用 detectron2 的样例。其一是读者大概率已经阅读过的 Colab Notebook ——骑马王子和气球检测，其二是使用命令行执行的 python 文件，包括演示文件 demo.py 及自行用于部署的 train_net.py & plain_train_net.py 。Notebook 已述明使用 Mask-RCNN 进行 mask detection 的简单步骤，包括注册数据集、配置 config 以及训练和验证、可视化结果。由于人力标注费力的原因，自行标注的数据集通常只有目标物体 box 而未标注 mask，因此需使用 Faster-RCNN 进行训练，有关代码可参考 Medium 文章 以及其 代码。

本篇参考 文章 ，简述 python 文件 + 命令行执行的项目部署方法，以便于在 GPU 服务器端运行。

目录一览

本项目部署目录如下：

Project

--configs

----COCO-Detection

faster_rcnn_R_50_FPN_3x.yaml # 官方代码库 copy

Base-RCNN-FPN.yaml # 同上

my_config.yaml # 自己的训练配置文件

--tools

train_net.py # 官方代码库 copy 并自行修改

train.sh

train_resume.sh

eval.sh

--utils

txt2coco.py # 将自己的数据集转化为标准 coco 数据集

数据集目录如下：

MyDataset

--train

01_00001.jpg # 图片名无所谓

...

--val

...

train.json # 位置自行决定

val.json

数据集注册

使用 CocoFormat 的数据集是最优雅的做法，当然你也可以按照官方给定的方法在 tools/rain_net.py 中自定义数据集。关于 csv/voc/labelme 等格式向 coco 格式的转化，可以参考 Github 的代码。本文 utils/txt2coco.py 基于上述代码将自行标注的 txt 格式数据转化为了 coco 格式。

为方便数据集注册，本文借鉴 文章 ，将数据集注册代码包装成类：

from detectron2.data import DatasetCatalog, MetadataCatalog

from detectron2.data.datasets.coco import load_coco_json

class Register:

"""用于注册自己的数据集"""

CLASS_NAMES = ['__background__', 'Red', 'Blue', 'Yellow', 'White', 'Black', 'Other', 'NoHelmet'] # 保留 background 类

ROOT = "/home/your_dataset_dir"

def __init__(self):

self.CLASS_NAMES = Register.CLASS_NAMES or ['__background__', ]

# 数据集路径

self.DATASET_ROOT = Register.ROOT or '/home/yourdir'

# ANN_ROOT = os.path.join(self.DATASET_ROOT, 'COCOformat')

self.ANN_ROOT = self.DATASET_ROOT

self.TRAIN_PATH = os.path.join(self.DATASET_ROOT, 'train')

self.VAL_PATH = os.path.join(self.DATASET_ROOT, 'val')

self.TRAIN_JSON = os.path.join(self.ANN_ROOT, 'train.json')

self.VAL_JSON = os.path.join(self.ANN_ROOT, 'val.json')

# VAL_JSON = os.path.join(self.ANN_ROOT, 'test.json')

# 声明数据集的子集

self.PREDEFINED_SPLITS_DATASET = {

"coco_my_train": (self.TRAIN_PATH, self.TRAIN_JSON),

"coco_my_val": (self.VAL_PATH, self.VAL_JSON),

}

def register_dataset(self):

"""

purpose: register all splits of datasets with PREDEFINED_SPLITS_DATASET

注册数据集(这一步就是将自定义数据集注册进Detectron2)

"""

for key, (image_root, json_file) in self.PREDEFINED_SPLITS_DATASET.items():

self.register_dataset_instances(name=key,

json_file=json_file,

image_root=image_root)

@staticmethod

def register_dataset_instances(name, json_file, image_root):

"""

purpose: register datasets to DatasetCatalog,

register metadata to MetadataCatalog and set attribute

注册数据集实例，加载数据集中的对象实例

"""

DatasetCatalog.register(name, lambda: load_coco_json(json_file, image_root, name))

MetadataCatalog.get(name).set(json_file=json_file,

image_root=image_root,

evaluator_type="coco")

def plain_register_dataset(self):

"""注册数据集和元数据"""

# 训练集

Da