使用Object Detection API训练自己的模型

前言

本次实验的环境是Mac high Sierra系统。

配置环境

这一部分参考的是models/installation.md at master · tensorflow/models · GitHub这个教程

• 首先安装Tensorflow和其依赖

# For CPU
pip install tensorflow
# For GPU
pip install tensorflow-gpu
# Dependencies
sudo apt-get install protobuf-compiler python-pil python-lxml python-tk
sudo pip install Cython
sudo pip install jupyter
sudo pip install matplotlib

• 编译Protobuf
  Tensorflow object detection API使用Protobuf来配置模型和训练的参数，所以需要先将Protobuf的库进行编译。在路径tensorflow/models/research/下执行

# From tensorflow/models/research/
protoc object_detection/protos/*.proto --python_out=.

• 向pythonpath添加库

# From tensorflow/models/research/
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

• 检查安装
  在resource路径下执行：

python object_detection/builders/model_builder_test.py

如果安装无误，会有如下结果：

Ran 15 tests in 0.082s

OK

学习object_detection_tutorial.ipynb

执行实例代码

在检查安装无误之后就可以尝试执行教程的ipython notebook，在执行部分代码块的时候可能会出现一些warning，可以暂时不管。
在下载模型的时候可能会出现下载失败的情况，需要翻墙。
执行结果如下：

加入自己的图片：

可以看到这个例子能够成功执行，由于下载的模型是速度快但是精度低的，所以识别的正确率还有待提升，有很多物品都没有识别到。

分析实例代码

这段代码的注释写得和详细，功能主要如下：

1. import所需的python库
2. 定义一些变量
3. 下载静态的模型并且加载到内存中
4. 加载一些辅助代码，例如将图片转成numpy中的数组
5. 图片识别的代码

使用API训练自己的模型

这个部分我参考的是Train your own object detector
我选取的检测对象是松鼠。

创建数据库

• 通过web scraper从https://pixabay.com 爬取图片链接。将爬取的图片链接导出到csv文件中。使用如下脚本将图片下载下来：

import csv
import os
import re
pattern = '(\d+)_(\d+)_(\d+).[jp][pn]g'
csv_reader = csv.reader(open('pixabay.csv', encoding='utf-8'))
for row in csv_reader:
    print(row[4])
    match = re.search(pattern,str(row[4]))
    if match:
    	name = match.group()
    	os.system("curl --output training/dataset/%s %s"%(name,row[4]))
    	print("OK")

• 安装和使用labelImg
  具体的安装过程可以参考Github上的readme文件https://github.com/tzutalin/labelImg。使用labelimg打开图片路径，然后用四个点标注出目标的位置，然后保存即可。
• 制作tfrecord文件
  先使用脚本将xml文件转换成csv文件，再将csv文件转换成tfrecord。

在本地进行训练

• 编辑label_map.pbtxt文件
  由于我的模型中只有一个类，所以此文件如下

item {
  id: 1
  name: 'squirrel'
}

• 使用官方教程中推荐的文件路径来存放训练时需要的各种文件：

+data
  -label_map file
  -train TFRecord file
  -eval TFRecord file
+models
  + model
    -pipeline config file
    +train
    +eval

• 编辑config文件
  关于文件的修改可以参考官方部分教程configuring。我直接使用的example\config文件中的ssd_mobilenet_v1_pets.config. 将此文件中的多个路径路径改为自己。教程中推荐从预训练的checkpoint开始训练，因为可以加快训练速度。可选择的预训练模型如下models/detection_model_zoo.md at master · tensorflow/models · GitHub 我使用了ssd_mobilenet_v1_coco_11_06_2017并将文件路径写在config文件的checkpoint路径里。
• 开始训练并查看Tensorboard
  从research文件路径下执行以下命令即可开始训练

python object_detection/train.py \
    --logtostderr \
    --pipeline_config_path=`pwd`/../../mytraining/models/model/pipeline.config \
    --train_dir=/Users/ver0n1ca/Desktop/training/mytraining/models/model/train

在research路径下执行一下命令可以对模型进行实时评估

python object_detection/eval.py \
    --logtostderr \
    --pipeline_config_path=`pwd`/../../mytraining/models/model/pipeline.config \
    --checkpoint_dir=/Users/ver0n1ca/Desktop/training/mytraining/models/model/train \
    --eval_dir=/Users/ver0n1ca/Desktop/training/mytraining/models/model/eval

执行以下命令可以可视化地显示训练结果

tensorboard --logdir=/Users/ver0n1ca/Desktop/training/mytraining/models/model/

训练结果

经过几个小时的训练之后，从tensorboard上可以看到classification lose降低到1左右：

localization loss降低到0.2以下

total loss为1.5左右

几个检测结果的实例如下：

本篇教程到这里就结束啦～