以Detectron用自己的数据集训练目标检测模型

在 2018-03-10 发布于 人工智能 下以来已有5,535人读过本文 | 0条评论 发表评论

前一段时间Facebook开源了其用于目标检测、识别和分割的框架——Detectron,该框架携带了这一领域里的经典和最新算法,诸如Faster-RCNN、RetinaNet、Mask-RCNN,等等。这绝对是造福这一领域从业者和研究者的事情,当然,也成了这一领域绕不开的东西了。

Detectron自带了可以基于Pascal VOC以及COCO数据集的模型训练方法。不过Detectron十分易用,即使训练自己的数据集,操作起来也十分简单方便。本文对此进行一下简单介绍和总结。

1.安装Detectron以及Caffe2

Detectron是基于Caffe2开发的,因此,首先需要安装Caffe2,然后从GitHub上clone下来Detectron进行安装,Detectron的安装说明文档十分详细且易于操作,直接参考进行安装即可。

1.1 Caffe2

安装完成Caffe2后,在Terminal里运行下面代码,检查是否安装成功:

# To check if Caffe2 build was successful
python -c 'from caffe2.python import core' 2>/dev/null && echo "Success" || echo "Failure"

# To check if Caffe2 GPU build was successful
# This must print a number > 0 in order to use Detectron
python -c 'from caffe2.python import workspace; print(workspace.NumCudaDevices())'

1.2 安装一些依赖及可能用到的COCO API

1.3 安装Detectron

从GitHub上将其clone下来,编译:

cd $DETECTRON/lib && make

然后检查是否安装成功:

python $DETECTRON/tests/test_spatial_narrow_as_op.py

2.配置自己的数据集

因为自定义自己的数据集比较繁琐,所以推荐使用Detectron自带的COCO数据集名称,并将自己的数据集转化为COCO数据集的格式,具体配置如下。

2.1 放置数据

将数据放置到$DETECTRON/lib/datasets/data下面,如animals数据集,并在其中以如下结构放置数据集的数据:

├── annos
├── annotations
├── classes.txt
└── images

其中,annos放置你的数据集的原始的标注文件,可能是txt,或者csv格式;classes.txt放置你标注的类别名称,每行一个类别,不含背景;images放置数据集的原始图像文件。annotations预备放置与COCO数据集格式的标注文件。下面需要将我们自己的标注文件生成COCO格式的标注文件。

2.2 转化标注格式

因为我们自己的数据集的标注文件是txt格式,每行一个标注,形似如下:

老虎 25 36 128 159
猫 25 36 128 159
狮子 25 36 128 159

所以,我们需要将这种格式的标注转换为COCO格式。程序如下:

#!/usr/bin/python
# -*- coding: utf-8 -*-
import json
import os
import sys

import cv2

if len(sys.argv) < 3:
  print "Usage: python convert_to_detectron_json.py root_path phase split"
  print "For example: python convert_to_detectron_json.py data train 100200"
  exit(1)

root_path = sys.argv[1]
phase = sys.argv[2]
split = int(sys.argv[3])

dataset = {
    'licenses': [],
    'info': {},
    'categories': [],
    'images': [],
    'annotations': []
}

with open(os.path.join(root_path, 'classes.txt')) as f:
  classes = f.read().strip().split()

for i, cls in enumerate(classes, 1):
  dataset['categories'].append({
      'id': i,
      'name': cls,
      'supercategory': 'beverage'
  })


def get_category_id(cls):
  for category in dataset['categories']:
    if category['name'] == cls:
      return category['id']


_indexes = sorted([f.split('.')[0] for f in os.listdir(os.path.join(root_path, 'annos'))])
if phase == 'train':
  indexes = [line for line in _indexes if int(line) > split]
else:
  indexes = [line for line in _indexes if int(line) <= split]

j = 1
for index in indexes:
  im = cv2.imread(os.path.join(root_path, 'images/') + index + '.jpg')
  height, width, _ = im.shape
  dataset['images'].append({
      'coco_url': '',
      'date_captured': '',
      'file_name': index + '.jpg',
      'flickr_url': '',
      'id': int(index),
      'license': 0,
      'width': width,
      'height': height
  })

  anno_file = os.path.join(root_path, 'annos/') + index + '.txt'
  with open(anno_file) as f:
    lines = [line for line in f.readlines() if line.strip()]

    for i, line in enumerate(lines):
      parts = line.strip().split()
      cls = parts[0]
      x1 = int(parts[1])
      y1 = int(parts[2])
      x2 = int(parts[3])
      y2 = int(parts[4])
      width = max(0, x2 - x1)
      height = max(0, y2 - y1)
      dataset['annotations'].append({
          'area': width * height,
          'bbox': [x1, y1, width, height],
          'category_id': get_category_id(cls),
          'id': j,
          'image_id': int(index),
          'iscrowd': 0,
          'segmentation': []
      })
      j += 1

folder = os.path.join(root_path, 'annotations')
if not os.path.exists(folder):
  os.makedirs(folder)
json_name = os.path.join(root_path, 'annotations/{}.json'.format(phase))
with open(json_name, 'w') as f:
  json.dump(dataset, f)

可以点击这里去GitHub查看源代码。你可以根据自己数据集的格式,改写程序,将标注转化为COCO格式。

2.3 设置指向自己的数据集

修改$DETECTRON/lib/datasets/dataset_catalog.py文件,将其中描述COCO数据集的两个部分修改为如下所示:

'coco_2014_train': {
        IM_DIR:
            _DATA_DIR + '/animals/images',
        ANN_FN:
            _DATA_DIR + '/animals/annotations/train.json'
},
'coco_2014_val': {
        IM_DIR:
            _DATA_DIR + '/animals/images',
        ANN_FN:
            _DATA_DIR + '/animals/annotations/val.json'
},

其中,_DATA_DIR是Detectron的数据的根目录,即$DETECTRON/lib/datasets/data,其余顾名思义。

3.新建配置文件,开始训练

到这里,设置好了相应的数据集,下面需要指定训练模型使用该数据集,Detectron是通过配置文件来指定的。

3.1 新建配置文件

可以直接复制Detectron自带的配置文件,并稍作修改即可,在$DETECTRON根目录下打开Terminal复制:

cd configs
mkdir animals
cp getting_started/tutorial_1gpu_e2e_faster_rcnn_R-50-FPN.yaml animals/animals_1gpu_e2e_faster_rcnn_R-50-FPN.yaml

然后打开animals_1gpu_e2e_faster_rcnn_R-50-FPN.yaml进行编辑:

MODEL:
  TYPE: generalized_rcnn
  CONV_BODY: FPN.add_fpn_ResNet50_conv5_body
  NUM_CLASSES: 5
  FASTER_RCNN: True
NUM_GPUS: 1
SOLVER:
  WEIGHT_DECAY: 0.0001
  LR_POLICY: steps_with_decay
  BASE_LR: 0.0025
  GAMMA: 0.1
  MAX_ITER: 60000
  STEPS: [0, 30000, 40000]
  # Equivalent schedules with...
  # 1 GPU:
  #   BASE_LR: 0.0025
  #   MAX_ITER: 60000
  #   STEPS: [0, 30000, 40000]
  # 2 GPUs:
  #   BASE_LR: 0.005
  #   MAX_ITER: 30000
  #   STEPS: [0, 15000, 20000]
  # 4 GPUs:
  #   BASE_LR: 0.01
  #   MAX_ITER: 15000
  #   STEPS: [0, 7500, 10000]
  # 8 GPUs:
  #   BASE_LR: 0.02
  #   MAX_ITER: 7500
  #   STEPS: [0, 3750, 5000]
RPN:
  ASPECT_RATIOS: (0.5, 1, 2)
FPN:
  FPN_ON: True
  MULTILEVEL_ROIS: True
  MULTILEVEL_RPN: True
  RPN_ASPECT_RATIOS: (0.5, 1, 2)
FAST_RCNN:
  ROI_BOX_HEAD: fast_rcnn_heads.add_roi_2mlp_head
  ROI_XFORM_METHOD: RoIAlign
  ROI_XFORM_RESOLUTION: 7
  ROI_XFORM_SAMPLING_RATIO: 2
TRAIN:
  WEIGHTS: pretrained_models/ImageNetPretrained/MSRA/R-50.pkl
  DATASETS: ('coco_2014_train',)
  SCALES: (500,)
  SNAPSHOT_ITERS: 5000
  MAX_SIZE: 833
  BATCH_SIZE_PER_IM: 256
  RPN_PRE_NMS_TOP_N: 2000  # Per FPN level
TEST:
  DATASETS: ('coco_2014_val',)
  SCALES: (500,)
  MAX_SIZE: 833
  NMS: 0.5
  RPN_PRE_NMS_TOP_N: 1000  # Per FPN level
  RPN_POST_NMS_TOP_N: 1000
  FORCE_JSON_DATASET_EVAL: True
OUTPUT_DIR: .

主要修改以下内容:

1.MODEL里的NUM_CLASSES,根据自己的数据集进行修改,包括背景。

2.TRAIN和TEST里的DATASETS,设置为如上所示,因为我们已经修改了dataset_catalog.py中的相应内容并指向了我们的数据集,所以我们可以直接使用这两个数据集的名字。

3.TRAIN里的WEIGHTS是一个预训练的模型,如果本地有,修改指向它即可。

4.关于NUM_GPUS,这个需要与SOLVER里面的学习率等一干参数进行联动设置。

3.2 开始训练

设置好数据集和配置文件,即可开始训练模型了。回到$DETECTRON根目录,打开Terminal运行如下命令即可开始训练:

python tools/train_net.py --cfg configs/animals/animals_1gpu_e2e_faster_rcnn_R-50-FPN.yaml OUTPUT_DIR detectron-output

这样即可使用Detectron基于自己的数据集训练各种模型了。

参考资料

Detectron官方设置数据集说明

Detectron官方GETTING STARTED

使用Detectron训练PASCAL VOC数据集

发表评论

您的昵称 *

您的邮箱 *

您的网站