TensorFlow 训练自己的目标检测器

TensorFlow 训练自己的目标检测器

        本文主要描述如何使用 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:/.../modes/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 自己训练好后的目标检测器都放在：

https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md

可以自己下载这些模型，一一查看检测效果。以下，假设你把某些预训练模型下载好了，放在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 train.py --train_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_dir /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文档