本文主要描述如何使用 Google 开源的目标检测 API 来训练目标检测器，内容包括：安装 TensorFlow/Object Detection API 和使用 TensorFlow/Object Detection API 训练自己的目标检测器。
一、安装 TensorFlow Object Detection API
Google 开源的目标检测项目 object_detection 位于与 tensorflow 独立的项目 models（独立指的是：在安装 tensorflow 的时候并没有安装 models 部分）内：models/research/object_detection。models 部分的 GitHub 主页为：
https://github.com/tensorflow/models
要使用 models 部分内的目标检测功能 object_detection，需要用户手动安装 object_detection。下面为详细的安装步骤：
1. 安装依赖项 matplotlib，pillow，lxml 等
使用 pip/pip3 直接安装：
$ sudo pip/pip3 install matplotlib pillow lxml
其中如果安装 lxml 不成功，可使用
$ sudo apt-get install python-lxml python3-lxml
安装。
2. 安装编译工具
$ sudo apt install protobuf-compiler
$ sudo apt-get install python-tk
$ sudo apt-get install python3-tk
3. 克隆 TensorFlow models 项目
使用 git 克隆 models 部分到本地，在终端输入指令：
$ git clone https://github.com/tensorflow/models.git
克隆完成后，会在终端当前目录出现 models 的文件夹。要使用 git（分布式版本控制系统），首先得安装 git：$ sudo apt-get install git。
4. 使用 protoc 编译
在 models/research 目录下的终端执行：
$ protoc object_detection/protos/*.proto --python_out=.
将 object_detection/protos/ 文件下的以 .proto 为后缀的文件编译为 .py 文件输出。
5. 配置环境变量
在 .bashrc 文件中加入环境变量。首先打开 .bashrc 文件：
$ sudo gedit ~/.bashrc
然后在文件末尾加入新行：
export PYTHONPATH=$PYTHONPATH:/.../models/research:/.../models/research/slim
其中省略号所在的两个目录需要填写为 models/research 文件夹、models/research/slim 文件夹的完整目录。保存之后执行如下指令：
$ source ~/.bashrc
让改动立即生效。
6. 测试是否安装成功
在 models/research 文件下执行：
$ python/python3 object_detection/builders/model_builder_test.py
如果返回 OK，表示安装成功。
二、训练 TensorFlow 目标检测器
成功安装好 TensorFlow Object Detection API 之后，就可以按照 models/research/object_detection 文件夹下的演示文件 object_detection_tutorial.ipynb 来查看 Google 自带的目标检测的检测效果。其中，Google 自己训练好后的目标检测器都放在：
可以自己下载这些模型，一一查看检测效果。以下，假设你把某些预训练模型下载好了，放在models/ research/ object_detection 的某个文件夹下，比如自定义文件夹 pretrained_models。
要训练自己的模型，除了使用 Google 自带的预训练模型之外，最关键的是需要准备自己的训练数据。
以下，详细列出训练过程（后续部分文章将详细介绍一些目标检测算法）：
1. 准备标注工具和文件格式转化工具
图像标注可以使用标注工具 labelImg，直接使用
$ sudo pip install labelImg
安装（当前好像只支持Python2.7）。另外，在此之前，需要安装它的依赖项 pyqt4：
$ sudo apt-get install pyqt4-dev-tools
（另一依赖项 lxml 前面已安装）。要使用 labelImg，只需要在终端输入 labelImg 即可。
为了方便后续数据格式转化，还需要准备两个文件格式转化工具：xml_to_csv.py 和 generate_tfrecord.py，它们的代码分别列举如下（它们可以从资料 [1] 中 GitHub 项目源代码链接中下载。其中为了方便一般化使用，我已经修改 generate_tfrecord.py 的部分内容使得可以自定义图像路径和输入 .csv 文件、输出 .record 文件路径，以及 6 中的 xxx_label_map.pbtxt 文件路径）：
(1) xml_to_csv.py 文件源码：
import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET
def xml_to_csv(path):
xml_list = []
for xml_file in glob.glob(path + '/*.xml'):
tree = ET.parse(xml_file)
root = tree.getroot()
for member in root.findall('object'):
value = (root.find('filename').text,
int(root.find('size')[0].text),
int(root.find('size')[1].text),
member[0].text,
int(member[4][0].text),
int(member[4][1].text),
int(member[4][2].text),
int(member[4][3].text)
)
xml_list.append(value)
column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
xml_df = pd.DataFrame(xml_list, columns=column_name)
return xml_df
def main():
image_path = os.path.join(os.getcwd(), 'annotations')
xml_df = xml_to_csv(image_path)
xml_df.to_csv('road_signs_labels.csv', index=None)
print('Successfully converted xml to csv.')
if __name__ == '__main__':
main()
(2) 修改后的 generate_tfrecord.py 文件源码：
"""
Usage:
# From tensorflow/models/
# Create train data:
python/python3 generate_tfrecord.py --csv_input=your path to read train.csv
--images_input=your path to read images
--output_path=your path to write train.record
--label_map_path=your path to read xxx_label_map.pbtxt
# Create validation data:
python/python3 generate_tfrecord.py --csv_input=you path to read val.csv
--images_input=you path to read images
--output_path=you path to write val.record
--label_map_path=your path to read xxx_label_map.pbtxt
"""
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
import os
import io
import pandas as pd
import tensorflow as tf
from PIL import Image
from object_detection.utils import dataset_util
from object_detection.utils import label_map_util
from collections import namedtuple
flags = tf.app.flags
flags.DEFINE_string('csv_input', '', 'Path to the CSV input')
flags.DEFINE_string('images_input', '', 'Path to the images input')
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
flags.DEFINE_string('label_map_path', '', 'Path to label map proto')
FLAGS = flags.FLAGS
def split(df, group):
data = namedtuple('data', ['filename', 'object'])
gb = df.groupby(group)
return [data(filename, gb.get_group(x)) for filename, x in
zip(gb.groups.keys(), gb.groups)]
def create_tf_example(group, label_map_dict, images_path):
with tf.gfile.GFile(os.path.join(
images_path, '{}'.format(group.filename)), 'rb') as fid:
encoded_jpg = fid.read()
encoded_jpg_io = io.BytesIO(encoded_jpg)
image = Image.open(encoded_jpg_io)
width, height = image.size
filename = group.filename.encode('utf8')
image_format = b'jpg'
xmins = []
xmaxs = []
ymins = []
ymaxs = []
classes_text = []
classes = []
for index, row in group.object.iterrows():
xmins.append(row['xmin'] / width)
xmaxs.append(row['xmax'] / width)
ymins.append(row['ymin'] / height)
ymaxs.append(row['ymax'] / height)
classes_text.append(row['class'].encode('utf8'))
classes.append(label_map_dict[row['class']])
tf_example = tf.train.Example(features=tf.train.Features(feature={
'image/height': dataset_util.int64_feature(height),
'image/width': dataset_util.int64_feature(width),
'image/filename': dataset_util.bytes_feature(filename),
'image/source_id': dataset_util.bytes_feature(filename),
'image/encoded': dataset_util.bytes_feature(encoded_jpg),
'image/format': dataset_util.bytes_feature(image_format),
'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
'image/object/class/label': dataset_util.int64_list_feature(classes),
}))
return tf_example
def main(_):
writer = tf.python_io.TFRecordWriter(FLAGS.output_path)
label_map_dict = label_map_util.get_label_map_dict(FLAGS.label_map_path)
images_path = FLAGS.images_input
examples = pd.read_csv(FLAGS.csv_input)
grouped = split(examples, 'filename')
for group in grouped:
tf_example = create_tf_example(group, label_map_dict, images_path)
writer.write(tf_example.SerializeToString())
writer.close()
output_path = FLAGS.output_path
print('Successfully created the TFRecords: {}'.format(output_path))
if __name__ == '__main__':
tf.app.run()
generate_tfrecord.py 也可以由 models/research/object_detection/dataset_tools 文件夹内的相关 .py 文件修改而来。后续也会有文章介绍怎么将图像转化为 .record 文件，敬请期待。
2. 创建工作目录，收集图片
在 Ubuntu 中新建项目文件夹，比如 xxx_detection（xxx 自取，下同），在该文件夹内新建文件夹 annotations，data，images，training。将所有收集到的图片放在 images 文件夹内。
3. 标注图片生成 xml 文件
利用标注工具 labelImg 对所有收集的图片进行标注，即将要检测的目标用矩形框框出，填入对应的目标类别名称，生成对应的 xml 文件，放在 annotations 文件夹内。
4. 将所有的 .xml 文件整合成 .csv 文件
执行 xml_to_csv.py（放在 xxx_detection文件夹下），将所有的 xml 标注文件汇合成一个 csv 文件，再从该 csv 文件中分出用于训练和验证的文件 train.csv 和 val.csv（分割比例自取），放入 data 文件夹。
5. 将 .csv 文件转化成 TensorFlow 要求的 .TFrecord 文件
将 generate_tfrecord.py 文件放在 TensorFlow models/research/object_detection 文件夹下，在该文件夹目录下的终端执行：
$ python3 generate_tfrecord.py --csv_input=/home/.../data/train.csv
--images_input=/home/.../images
--output_path=/home/.../data/train.record
--label_map_path=/home/.../training/xxx_label_map.pbtxt
类似的，对 val.csv 执行相同操作，生成 val.record 文件。（其中 xxx_label_map.pbtxt 文件见下面的 6）
6. 编写 .pbtxt 文件
仿照 TensorFlow models/research/object_detection/data 文件夹下的 .pbtxt 文件编写自己的 .pbtxt 文件：对每个要检测的类别写入
item {
id: k
name: ‘xxx’
}
其中 item 之间空一行，类标号从 1 开始，即 k >= 1。将 .pbtxt 文件命名为 xxx_label_map.pbtxt 并放入training 文件夹。
7. 配置 .config 文件
从 TensorFlow models/research/object_detection/samples/configs 文件夹内选择合适的一个 .config 文件复制到项目工程的 training 文件夹内，将名称改为与工程相关的 保留模型名 _xxx.config（其中保留模型名为原 .config 文件关于模型的命名字段，建议命名时保留下来，xxx 为与项目相关的自己命名字段），打开文件作如下修改：
（1）修改模型参数
将 model {} 中的 num_classes 修改为工程要检测的类别个数。另外，也可以修改训练参数：
train_config: {} => num_steps: xxx => schedule {} => step = xxx
num_steps 表示将要训练的次数，删除这一行为不确定次数训练（随时可用 Ctrl+C 中断），后面的 step 表示学习率每过 step 步后进行衰减。这些参数由自己的经验确定，也可以使用默认值。
其它参数一般不需要修改。
（2）修改文件路径
将 .config 文件中所有的 ’PATH_TO_BE_CONFIGURED’ 文件路径修改为相应的 .ckpt（预训练模型文件路径），.record，.pbtxt 文件所在路径。
将修改后的 保留模型名_xxx.config 文件放在 training 文件夹内。
8. 开始本地训练目标检测器
在 TensorFlow models/research/object_detection 目录下的终端执行：
$ python3 model_main.py --model_dir=/home/.../training
--pipeline_config_path=/home/.../training/保留模型名_xxx.config
进行模型训练，期间每隔一定时间会输出若干文件到 training 文件夹。在训练过程中可使用 Ctrl+C 任意时刻中断训练，之后再执行上述代码会从断点之处继续训练，而不是从头开始（除非把训练输出文件全部删除）。
9. 查看实时训练曲线
在任意目录下执行：
$ tensorboard --logdir=/home/.../training
打开返回的 http 链接查看 Loss 等曲线的实时变化情况。
10. 导出 .pb 文件用于推断
模型训练完后，生成的 .ckpt 文件已经可以调用进行目标检测。也可以将 .ckpt 文件转化为 .pb 文件用于推断。在 TensorFlow models/research/object_detection 目录下的终端执行：
$ python3 export_inference_graph.py --input_type image_tensor
--pipeline_config_path /home/.../training/pipeline.config
--trained_checkpoint_prefix /home/.../training/model.ckpt-200000
--output_directory /home/.../training/output_inference_graph
执行上述代码之后会在 /home/.../training 文件夹内看到新的文件夹 output_inference_graph，里面存储着训练好的最终模型，如直接调用的用于推断的文件：frozen_inference_graph.pb。其中命令中 model.ckpt-200000 表示训练 200000 生成的模型，实际执行上述代码时要修改为自己训练多少次后生成的模型。其它路径和文件（夹）名称也由自己任意指定。
11. 调用训练好的模型进行目标检测
调用 frozen_inference_graph.pb 进行目标检测请参考 TensorFlow models/research/object_detection 文件夹下的 object_detection_tutorial.ipynb 。但该文件只针对单张图像，对多张图像不友好，因为每检测一张图像都要重新打开一个会话（语句 with tf.Session() as sess每张图像执行一次），而这是非常耗时的操作。可以改成如下的形式：
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Nov 4 15:05:09 2017
@author: shirhe-lyh
"""
import time
import cv2
import numpy as np
import tensorflow as tf
#--------------Model preparation----------------
# Path to frozen detection graph. This is the actual model that is used for
# the object detection.
PATH_TO_CKPT = 'path_to_your_frozen_inference_graph.pb'
# Load a (frozen) Tensorflow model into memory
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
# Each box represents a part of the image where a particular
# object was detected.
gboxes = detection_graph.get_tensor_by_name('detection_boxes:0')
# Each score represent how level of confidence for each of the objects.
# Score is shown on the result image, together with the class label.
gscores = detection_graph.get_tensor_by_name('detection_scores:0')
gclasses = detection_graph.get_tensor_by_name('detection_classes:0')
gnum_detections = detection_graph.get_tensor_by_name('num_detections:0')
# TODO: Add class names showing in the image
def detect_image_objects(image, sess, detection_graph):
# Expand dimensions since the model expects images to have
# shape: [1, None, None, 3]
image_np_expanded = np.expand_dims(image, axis=0)
# Actual detection.
(boxes, scores, classes, num_detections) = sess.run(
[gboxes, gscores, gclasses, gnum_detections],
feed_dict={image_tensor: image_np_expanded})
# Visualization of the results of a detection.
boxes = np.squeeze(boxes)
scores = np.squeeze(scores)
height, width = image.shape[:2]
for i in range(boxes.shape[0]):
if (scores is None or
scores[i] > 0.5):
ymin, xmin, ymax, xmax = boxes[i]
ymin = int(ymin * height)
ymax = int(ymax * height)
xmin = int(xmin * width)
xmax = int(xmax * width)
score = None if scores is None else scores[i]
font = cv2.FONT_HERSHEY_SIMPLEX
text_x = np.max((0, xmin - 10))
text_y = np.max((0, ymin - 10))
cv2.putText(image, 'Detection score: ' + str(score),
(text_x, text_y), font, 0.4, (0, 255, 0))
cv2.rectangle(image, (xmin, ymin), (xmax, ymax),
(0, 255, 0), 2)
return image
with detection_graph.as_default():
with tf.Session(graph=detection_graph) as sess:
video_path = 'path_to_your_video'
capture = cv2.VideoCapture(video_path)
while capture.isOpened():
if cv2.waitKey(30) & 0xFF == ord('q'):
break
ret, frame = capture.read()
if not ret:
break
t_start = time.clock()
detect_image_objects(frame, sess, detection_graph)
t_end = time.clock()
print('detect time per frame: ', t_end - t_start)
cv2.imshow('detected', frame)
capture.release()
cv2.destroyAllWindows()
这样改动之后，有好处也有坏处，好处是处理视频或很多图像时只生成一次会话节省时间，而且从原文件中去掉了语句：
sys.path.append("..")
from object_detection.utils import ops as utils_ops
from utils import label_map_util
from utils import visualization_utils as vis_util
使得在任意目录下都可以执行。坏处是：上述代码没有使用 label_map_util 和 vis_util 等这些 object_detection 伴随的模块，使得检测结果显示的时候只能自己利用 OpenCV 来做，而存在一个较大的缺陷：不能显示检测出的目标的类别名称（待完善）。
资料：
[1]目标干脆面君：动动手，用TensorFlow训练自己的目标检测模型，36kr
[2]利用TensorFlow Object Detection API训练自己的数据集，红黑联盟
[3]TensorFlow models/research/object_detection的GitHub文档
