易 AI - 使用 TensorFlow Object Detection API 训练自定义目标检测模型

ai deep learning machine learning 人工智能 机器学习 深度学习 tensorflow object detection 目标检测 tensorflow object detection api r-cnn ssd labelimg labelme

Read in English

前言

本文将介绍目标检测（Object Detection）的概念，并通过案例讲解如何使用 TensorFlow Object Detection API 来训练自定义的目标检测器，包括：数据集采集和制作、TensorFlow Object Detection API 安装以及模型的训练。

案例效果如下图所示：

目标检测

如上图所示，图像分类解决的问题是图中的物体是什么，而目标检测能识别图片中有哪些物体以及物体的位置（坐标）。

位置

目标检测的位置信息一般有两种格式：

  • 极坐标(xmin, ymin, xmax, ymax)：
    • xmin,ymin：x,y 坐标的最小值
    • xmin,ymin：x,y 坐标的最大值
  • 中心点：(x_center, y_center, w, h)
    • x_center, y_center：目标检测框的中心点坐标
    • w,h：目标检测框的宽、高

注：图片左上角为原点(0,0)

发展史

传统方法（候选区域+手工特征提取+分类器）

HOG+SVM、DPM

Region Proposal+CNN（Two-stage）

R-CNN, SPP-NET, Fast R-CNN, Faster R-CNN

端到端（One-stage）

YOLO、SSD

TensorFlow Object Detection API

TensorFlow Object Detection API 是一个构建在 TensorFlow 之上的开源框架，可以轻松构建、训练和部署对象检测模型。另外，TensorFlow Object Detection API 还提供了 Model Zoo 方便我们选择和切换预训练模型。

安装依赖项

使用以下命令检查是否安装成功。

1
2
3
4

$ conda --version
conda 4.9.2
$ protoc --version
libprotoc 3.17.1

安装 API

TensorFlow Object Detection API 官方的安装步骤较为繁琐，笔者写了一个脚本直接一键安装。

执行 git clone https://github.com/CatchZeng/object-detection-api.git 下载仓库，然后到该仓库（下文简称 oda 仓库）目录下，执行以下命令，如果看到如下输出，表示安装成功。

1
2
3
4
5
6

$ conda create -n  od python=3.8.5 && conda activate od && make install
......
----------------------------------------------------------------------
Ran 24 tests in 21.869s

OK (skipped=1)

注：如果你不想用 conda，可以在自己的 python 环境上直接使用 make install 安装即可，比如在 colab 中使用。

注：由于 cudaDNN 和 toolkit 更新可能没有 TensorFlow 快。因此，如果你的机器有 GPU，安装完成后，需要将 TensorFlow 降回到 cudaDNN 和 toolkit 支持的版本这样才能支持 GPU 训练，以 2.8.0 为例：

1
2

$ pip install --upgrade tf-models-official==2.8.0
$ pip install --upgrade tensorflow==2.8.0

注：如果安装失败，可以参考官方文档的详细步骤安装。

工程创建

注：!!! 从这里开始，请确保在 conda od 的环境下执行。

1
2
3
4
5
6
7

$ conda activate od
$ conda env list
# conda environments:
#
od                    *  /Users/catchzeng/.conda/envs/od
tensorflow               /Users/catchzeng/.conda/envs/tensorflow
base                     /Users/catchzeng/miniconda3

到 oda 仓库目录下，执行以下命令，创建工程目录结构。

注：SAVE_DIR 为保存项目的目录，NAME 为项目的名称。

1

$ make workspace-box SAVE_DIR=workspace NAME=test

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

└── workspace
    └── test
        ├── Makefile
        ├── annotations：存放标注好的数据集数据（val.record、train.record、label_map.pbtxt）
        ├── convert_quant_lite.py：量化 tflite 模型脚本
        ├── export_tflite_graph_tf2.py：导出 tflite 模型脚本
        ├── exported-models：存放训练完之后导出的模型
        ├── exporter_main_v2.py：导出模型脚本
        ├── images：数据集图片和 xml 标注
        │   ├── test：手动验证图片
        │   ├── train：训练集图片和 xml 标注
        │   └── val：验证集图片和 xml 标注
        ├── model_main_tf2.py：训练模型脚本
        ├── models：自定义模型
        ├── pre-trained-models：TensorFlow Model Zoo 提供的预训练模型
        └── test_images.py：手动验证图片脚本

数据集

图片

笔者喜欢喝茶，这次就用茶杯(cup)、茶壶(teapot)、加湿器(humidifier) 来做案例吧。

将收集的图片，放入工程目录的 images 的三个子目录下。

注：本案例只是为了验证如何训练目标识别模型，因此数据集采集得比较少，实际项目中记得尽量采集多点数据集。

标注

收集完图片后，需要对训练和验证集图片进行标注。标注工具，选用较为常用的 LabelImg。

根据 installation 的说明安装好 LabelImg，然后执行 labelImg 选择 train 和 val 文件夹进行标注。

标注完成后，会生成图片对应的 xml 标注文件，如下所示：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

workspace/test/images
├── test
│   ├── 15.jpg
│   └── 16.jpg
├── train
│   ├── 1.jpg
│   ├── 1.xml
│   ├── 10.jpg
│   ├── 10.xml
│   ├── 2.jpg
│   ├── 2.xml
│   ├── 3.jpg
│   ├── 3.xml
│   ├── 4.jpg
│   ├── 4.xml
│   ├── 5.jpg
│   ├── 5.xml
│   ├── 6.jpg
│   ├── 6.xml
│   ├── 7.jpg
│   ├── 7.xml
│   ├── 8.jpg
│   ├── 8.xml
│   ├── 9.jpg
│   └── 9.xml
└── val
    ├── 11.jpg
    ├── 11.xml
    ├── 12.jpg
    ├── 12.xml
    ├── 13.jpg
    ├── 13.xml
    ├── 14.jpg
    └── 14.xml

创建 TFRecord

TensorFlow Object Detection API 只支持 TFRecord 格式的数据集，因此，需要把标注好的数据集进行转换。

先 cd 到工程目录（cd workspace/test），然后执行 make gen-tfrecord，将在 annotations 文件夹下生成 label_map.pbtxt 和 TFRecord 格式的数据集。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

$ make gen-tfrecord
python gen_label_map.py
unsorted:  ['cup', 'teapot', 'humidifier']
sorted:  ['cup', 'humidifier', 'teapot']
item {
  id: 1
  name: 'cup'
}

item {
  id: 2
  name: 'humidifier'
}

item {
  id: 3
  name: 'teapot'
}

python generate_tfrecord.py \
        -x images/train \
        -l annotations/label_map.pbtxt \
        -o annotations/train.record
Successfully created the TFRecord file: annotations/train.record
python generate_tfrecord.py \
        -x images/val \
        -l annotations/label_map.pbtxt \
        -o annotations/val.record
Successfully created the TFRecord file: annotations/val.record

1
2
3
4

annotations
├── label_map.pbtxt
├── train.record
└── val.record

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

# label_map.pbtxt
item {
  id: 1
  name: 'cup'
}

item {
  id: 2
  name: 'humidifier'
}

item {
  id: 3
  name: 'teapot'
}

模型训练

注：!!! 从这里开始，请确保已经 cd 到工程目录（cd workspace/test）。

下载预训练模型

从 Model Zoo 中选择合适的模型下载解压并放到 workspace/test/pre-trained-models 中。

如果你选择的是 SSD MobileNet V2 FPNLite 320x320 可以执行如下命令，自动下载并解压

1

$ make dl-model

目录结构如下：

1
2
3
4
5
6

└── test
    └── pre-trained-models
        └── ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8
            ├── checkpoint
            ├── pipeline.config
            └── saved_model

配置训练 Pipeline

在 models 目录创建对应的模型文件夹，比如：ssd_mobilenet_v2_fpnlite_320x320，并拷贝 pre-trained-models/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8/pipeline.config。

1
2
3
4
5

└── test
    ├── models
    │   └── ssd_mobilenet_v2_fpnlite_320x320
    │       └── pipeline.config
    └── pre-trained-models

其中，pipeline.config 如下几处需要根据项目修改

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

model {
  ssd {
    num_classes: 3 # 修改为需要识别的目标个数，示例项目为 3 种
    ......
}
train_config {
  batch_size: 8 # 这里需要根据自己的配置，调整大小，这里设置为 8
  ......
  optimizer {
    momentum_optimizer {
      learning_rate {
        cosine_decay_learning_rate {
          learning_rate_base: 0.07999999821186066
          total_steps: 10000 # 修改为想要训练的总步数
          warmup_learning_rate: 0.026666000485420227
          warmup_steps: 1000
        }
      }
      momentum_optimizer_value: 0.8999999761581421
    }
    use_moving_average: false
  }
  fine_tune_checkpoint: "pre-trained-models/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8/checkpoint/ckpt-0" # 修改为预制模型的路径
  num_steps: 10000 # 修改为想要训练的总步数
  startup_delay_steps: 0.0
  replicas_to_aggregate: 8
  max_number_of_boxes: 100
  unpad_groundtruth_tensors: false
  fine_tune_checkpoint_type: "detection" # 这里需要修改为 detection，因为我们是做目标检测
  fine_tune_checkpoint_version: V2
}
train_input_reader {
  label_map_path: "annotations/label_map.pbtxt" # 修改为标注的路径
  tf_record_input_reader {
    input_path: "annotations/train.record" # 修改为训练集的路径
  }
}
eval_config {
  metrics_set: "coco_detection_metrics"
  use_moving_averages: false
}
eval_input_reader {
  label_map_path: "annotations/label_map.pbtxt" # 修改为标注的路径
  shuffle: false
  num_epochs: 1
  tf_record_input_reader {
    input_path: "annotations/val.record" # 修改为验证集的路径
  }
}

训练模型

1

$ make train

注：如遇以下问题

ValueError: numpy.ndarray size changed, may indicate binary incompatibility. Expected 88 from C header, got 80 from PyObject

执行以下命令，重新安装下 numpy 即可。

1
2

pip uninstall numpy
pip install numpy

模型导出与转换

  • 普通模型

    1

    $ make export

  • TFLite 模型

    1

    $ make export-lite

  • 转换 TFLite 模型

    1

    $ make convert-lite

  • 转换量化版 TFLite 模型

    1

    $ make convert-quant-lite

注：以上命令，大家可以加 -640 表示使用 SSD MobileNet V2 FPNLite 640x640 的模型，比如：make train -> make train-640

测试模型

执行 make export 导出模型后，将测试图片放到 images/test 文件夹下，然后执行 python test_images.py 即可输出标记好目标的图片到 images/test_annotated。

小结

本文通过案例将目标检测的整个流程都过了一遍，希望能帮助大家快速掌握训练自定义目标检测器的能力。

案例的代码和数据集都已经放在了 https://github.com/CatchZeng/object-detection-api，有需要的同学可以自行获取。

后面的文章将会为大家带来，目标检测的原理、常用的目标检测网络，以及目标分割。本篇就到这了，咱们下一篇见。

延伸阅读

参考链接


CatchZeng
Written by CatchZeng Follow
AI (Machine Learning) and DevOps enthusiast.
Google Optimize 变体的移动端编辑与 Shopify 不兼容
The easiest way to Train a Custom Object Detection Model Using TensorFlow Object Detection API