TensorFlow Object Detection API使用小结

本文所用的 Python 版本为 python-3.6.2，TensorFlow 版本为tensorflow-1.4.0，编程语言为 python3，系统环境为 Windows 10。

前言

　　很久没写过东西了，主要原因是最近研究生课程开始陆续结课，Shaun 也要忙于应付各种结课时的考试、论文、project 等一堆麻烦事。这不深度学习结课时需要做个 project，Shaun 也顺便将做这个 project 的过程记录下来。

准备篇

　　该 project 主要利用 TensorFlow 中的 Object Detection API 进行训练和检测。在开始使用该 API 之前需要安装配置 Python 环境。

　　既然是 Python 首先需要 下载安装Python，安装完之后，为了顺利使用 pip 需要配置环境变量，在 Windows 系统环境变量中 Path 末尾添加：

变量名	变量值
Path	;C:\Users\admin\AppData\Local\Programs\Python\Python36\; C:\Users\admin\AppData\Local\Programs\Python\Python36\Scripts\

其中 C:\Users\admin\AppData\Local\Programs\Python\Python36 为 python-3.6.2 默认安装目录。

　　然后为了方便使用命令行工具，下载安装git，安装方式一路默认即可。

　　接下来利用 pip 安装 TensorFlow，鼠标右键桌面空白处，点击“Git Bash Here”，打开 bash 命令行，输入 pip install tensorflow，其中一些依赖关系可能需要手动解决，手动解决的办法就是用 pip install 相关依赖库，这是 CPU 版的 TensorFlow，若要使用 GPU，则需要安装 GPU 版的 TensorFlow，安装命令为：pip install tensorflow-gpu，以同样方式解决依赖关系。由于 Shaun 电脑没 N 卡，所以没安装 GPU 版的 TensorFlow，所以如果想使用 GPU 版的 TensorFlow 请另行尝试。

　　然后安装 TensorFlow Object Detection API 依赖库，在命令行中输入：

pip install pillow
pip install lxml
pip install jupyter
pip install matplotlib 

　　因为 tensorflow 并没有默认自带 Object Detection API，所以该 API 需要自行下载，下载地址为：https://github.com/tensorflow/models ，下载之后解压，Shaun 解压目录为：D:\ProgramFiles\PythonLibs\tensorflow，解压完之后需要配置环境目录，在系统环境目录中添加：

变量名	变量值
PYTHONPATH	D:\ProgramFiles\PythonLibs\tensorflow\models; D:\ProgramFiles\PythonLibs\tensorflow\models\research; D:\ProgramFiles\PythonLibs\tensorflow\models\research\slim;

　　下载配置 Object Detection API 完之后需要安装 Protoc，进入 Protoc下载页，下载 protoc-3.4.0-win32.zip，解压之后将 bin 文件夹内的 protoc.exe 拷贝到 C:\windows\system32 目录下（用于将 protoc.exe 所在的目录配置到环境变量当中）,当然也可以在系统环境变量 Path 中添加该 bin 文件夹路径。

　　最后在命令行中切换目录至：D:\ProgramFiles\PythonLibs\tensorflow\models\research 文件夹，即 object_detection 文件夹所在目录，在命令行中输入：

protoc object_detection/protos/*.proto --python_out=.

编译 object_detection/protos 文件夹下的 proto 文件，生成对应的 python 文件。

　　至此，Windows 下 TensorFlow中 的 Object Detection API 的使用配置全部完成，至于 Ubuntu 下的配置可参考其官方文档。

　　至于如何验证，可以在命令行中切换目录至 object_detection，输入：jupyter notebook，稍等一会，浏览器将自动打开 http://localhost:8888/tree jupyter 界面，点击 object_detection_tutorial.ipynb 文件，进入打开的新标签，点击“Cell”中的“Run All”，耐心等待几 ~ 十几分钟（因为它需要下载相应的模型），将会在浏览器下方显示检测结果。

　　截止本文完成前，该API公开的有以下几个模型：

Model name	Speed (ms)	COCO mAP1	Outputs
ssd_mobilenet_v1_coco	30	21	Boxes
ssd_inception_v2_coco	42	24	Boxes
faster_rcnn_inception_v2_coco	58	28	Boxes
faster_rcnn_resnet50_coco	89	30	Boxes
faster_rcnn_resnet50_lowproposals_coco	64		Boxes
rfcn_resnet101_coco	92	30	Boxes
faster_rcnn_resnet101_coco	106	32	Boxes
faster_rcnn_resnet101_lowproposals_coco	82		Boxes
faster_rcnn_inception_resnet_v2_atrous_coco	620	37	Boxes
faster_rcnn_inception_resnet_v2_atrous_lowproposals_coco	241		Boxes
faster_rcnn_nas	1833	43	Boxes
faster_rcnn_nas_lowproposals_coco	540		Boxes

　　根据上述模型可推知，利用该 API 可能只能训练 Faster-RCNN、R-FCN 和 SSD 三种算法的模型。

接下来介绍如何使用该 API 来训练自己的模型进行物体检测。

实践篇

数据准备篇

　　既然要训练自己的模型，当然要准备相应的数据，而 TensorFlow 有其独特的输入数据格式 TFRecord，所以通常还要将自己的数据转换成 TFRecord 格式以输入 TensorFlow 中进行训练。以 datitran/raccoon_dataset 数据集为例，该作者在 Google image 上收集了 200 张 Raccoon 图片，用 LabelImg 对这些图片进行标记，并将标记以 PASCAL VOC 格式保存为 xml 文件。作者在文中也提到了另一个图片标记工具 FIAT (Fast Image Data Annotation Tool) 。保存为 PASCAL VOC 格式的 xml 文件之后，可以使用 object_detection 文件夹中的 create_pascal_tf_record.py 文件将数据转化为 TFRecord 格式，用法为：

./create_pascal_tf_record --data_dir=/home/user/VOCdevkit \
        --year=VOC2012 \
        --output_path=/home/user/pascal.record

当然也可以使用 datitran 作者提供的 xml_to_csv.py 文件将 xml 文件先转化为 csv 文件，再利用 generate_tfrecord.py 文件将 csv 文件转化成 TFRecord 格式文件。

　　注意，使用 xml_to_csv.py 和 generate_tfrecord.py 其文件结构应该是这样的：

.
├── annotations
├── generate_tfrecord.py
├── images
└── xml_to_csv.py

2 directories, 2 files

其中 images 文件夹存的是 jpg 图片，annotations 文件夹存的是 xml 标签文件。generate_tfrecord.py 文件中的：

def class_text_to_int(row_label):
    if row_label == 'raccoon':
        return 1
    else:
None

其中的 raccoon 注意要改成自己的类别标签。如此，数据的问题就解决了。

训练篇

　　然后就是正式开始训练了，以 Faster-RCNN 算法为例。首先准备相应的数据，Shaun 准备的数据文件结构如下：

TensorFlowObjectDetectionAPITest
├── data
│ 　├── model.ckpt.data-00000-of-00001
│ 　├── model.ckpt.index
│ 　├── model.ckpt.meta
│ 　├── object_label_map.pbtxt
│ 　├── test.record
│ 　└── train.record
├── eval
├── eval.py
├── export_inference_graph.py
├── faster_rcnn_resnet101_coco.config
├── model
├── train
└── train.py

4 directories, 10 files

其中，TensorFlowObjectDetectionAPITest 为项目文件夹，该 project 在此文件夹下运行；

data 文件夹中三个 model.ckpt 文件：model.ckpt.data-00000-of-00001、model.ckpt.index 和 model.ckpt.meta 来自 faster_rcnn_resnet101_coco 模型，用来初始化网络参数；

object_label_map.pbtxt 文件内容如下：

item { ​ id: 1 ​ name: 'raccoon' }

将其中的 raccoon 改成自己的类别标签，如果有多个类别标签则可以参考 object_detection\data 文件夹中的 pascal_label_map.pbtxt 文件格式；

test.record 和 train.record 是生成的 TFRecord 数据，分别为待输入的测试数据和训练数据；

eval 文件夹为空文件夹用来输出测试结果；train 文件夹为空文件夹用来输出训练结果（包括checkpoint文件和最终的模型文件）；

faster_rcnn_resnet101_coco.config 为配置文件，包括各种参数和输入输出数据的配置，其来自 object_detection\samples\configs 文件夹中 faster_rcnn_resnet101_coco.config 文件，在使用时需对其做如下修改：

1. 首先是 num_classes，这是待检测的类别数目，如果只要检测一种，则将其值改为 1；

2. fine_tune_checkpoint: "PATH_TO_BE_CONFIGURED/model.ckpt"，将 PATH_TO_BE_CONFIGURED 改为 ./data；

3. ``` train_input_reader: { tf_record_input_reader { input_path: "PATH_TO_BE_CONFIGURED/mscoco_train.record" } label_map_path: "PATH_TO_BE_CONFIGURED/mscoco_label_map.pbtxt" }

      将其中的的 `PATH_TO_BE_CONFIGURED/mscoco_train.record` 改为 `./data/train.record`，将其中的 `PATH_TO_BE_CONFIGURED/mscoco_label_map.pbtxt` 改为 `./data/object_label_map.pbtxt`；
   
   4. ```
      eval_input_reader: {
        tf_record_input_reader {
          input_path: "PATH_TO_BE_CONFIGURED/mscoco_val.record"
        }
        label_map_path: "PATH_TO_BE_CONFIGURED/mscoco_label_map.pbtxt"
        shuffle: false
        num_readers: 1
        num_epochs: 1
      }

   将其中的的 PATH_TO_BE_CONFIGURED/mscoco_val.record 改为 ./data/test.record，将其中的 PATH_TO_BE_CONFIGURED/mscoco_label_map.pbtxt 改为 ./data/object_label_map.pbtxt；

至于其它的参数可以选择默认，不对其进行修改；

train.py 为训练代码，其来自 object_detection/ 文件夹中的 train.py，直接复制出来使用即可，具体用法为：

python train.py --logtostderr --train_dir=./train --pipeline_config_path=faster_rcnn_resnet101_coco.config

其在运行过程中会在 train 文件夹生成一系列训练过程文件，比如 checkpoint、model.ckpt-{num}（{num} 代表训练过程保存的第几个网络模型，一般从 0 开始，包括 .index、.meta和 .data 三个文件）等文件。

eval.py 为评估代码，其来自 object_detection/ 文件夹中的 eval.py，直接复制出来使用即可，具体用法为：

python eval.py --logtostderr --checkpoint_dir=./train --eval_dir=./eval --pipeline_config_path=./faster_rcnn_resnet101_coco.config

其在运行过程中会在 eval 文件夹生成一系列评估文件，每个文件对应一个测试 image。

export_inference_graph.py 为导出 pb 模型代码，其来自 object_detection/ 文件夹中的 export_inference_graph.py，直接复制出来使用即可，具体用法为：

python export_inference_graph.py --input_type image_tensor --pipeline_config_path ./faster_rcnn_resnet101_coco.config --trained_checkpoint_prefix ./train/model.ckpt-18298 --output_directory ./model

其中 model.ckpt-18298 表示使用第 18298 次保存的网络模型导出 pb 模型文件，导出的模型文件保存在 model 文件夹，主要有一下几个文件：

- graph.pbtxt

- model.ckpt.data-00000-of-00001

- model.ckpt.info

- model.ckpt.meta

- frozen_inference_graph.pb

其中 frozen_inference_graph.pb 就是训练成功用来检测目标的模型。

　　TensorFlow 训练时可以随时查看训练过程，如损失函数的值下降曲线等，所用命令为：在命令行中切换目录至 project 运行目录，即 train.py 所在文件夹，Shaun 这里即 TensorFlowObjectDetectionAPITest 文件夹，输入：tensorboard --logdir=./，等待片刻，在浏览器地址栏输入：http://localhost:6006/，即可看到训练过程曲线。

检测篇

　　检测结果使用 opencv 窗口显示（至于 python 中 opencv 的使用详见下一篇（PyQt5使用小结）），具体调用自己训练的模型进行检测的 Python 代码（该代码为 eli 大佬参考 object_detection 文件夹中的 object_detection_tutorial.ipynb（该文件可在 jupyter 中查看）改的）为：

import cv2
import numpy as np
import tensorflow as tf
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as vis_util


class Detector(object):
    def __init__(self):
        self.PATH_TO_CKPT = r'./model/frozen_inference_graph.pb'    # 选择模型
        self.PATH_TO_LABELS = r'./data/object_label_map.pbtxt'      # 选择类别标签文件
        self.NUM_CLASSES = 1
        self.detection_graph = self._load_model()
        self.category_index = self._load_label_map()

    def _load_model(self):
        detection_graph = tf.Graph()
        with detection_graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(self.PATH_TO_CKPT, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')
        return detection_graph

    def _load_label_map(self):
        label_map = label_map_util.load_labelmap(self.PATH_TO_LABELS)
        categories = label_map_util.convert_label_map_to_categories(label_map,
                                                                    max_num_classes=self.NUM_CLASSES,
                                                                    use_display_name=True)
        category_index = label_map_util.create_category_index(categories)
        return category_index

    def detect(self, image):
        with self.detection_graph.as_default():
            with tf.Session(graph=self.detection_graph) as sess:
                # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
                image_np_expanded = np.expand_dims(image, axis=0)
                image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
                boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
                scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
                classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
                num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')
                # Actual detection.
                (boxes, scores, classes, num_detections) = sess.run(
                    [boxes, scores, classes, num_detections],
                    feed_dict={image_tensor: image_np_expanded})
                # Visualization of the results of a detection.
                vis_util.visualize_boxes_and_labels_on_image_array(
                    image,
                    np.squeeze(boxes),
                    np.squeeze(classes).astype(np.int32),
                    np.squeeze(scores),
                    self.category_index,
                    use_normalized_coordinates=True,
                    line_thickness=8)

        cv2.namedWindow("detection", cv2.WINDOW_NORMAL)
        cv2.imshow("detection", image)
        cv2.waitKey(0)

if __name__ == '__main__':
    image = cv2.imread('./test_img.jpg') # 选择待检测的图片
    detector = Detector()
    detector.detect(image)

后记

　　经过这次 TensorFlow 训练，感觉深度学习 真tm 吃硬件，费时间，也难怪神经网络理论出来几十年之后才火，当年的硬件根本无法支持这么大的计算量。

附录

最后附上 datitran 作者提供的 xml_to_csv.py 文件源码和 generate_tfrecord.py 文件源码：

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('raccoon_labels.csv', index=None)
    print('Successfully converted xml to csv.')


main()

generate_tfrecord.py 文件源码 如下：

"""
Usage:
  # From tensorflow/models/
  # Create train data:
  python generate_tfrecord.py --csv_input=data/train_labels.csv  --output_path=train.record
  # Create test data:
  python generate_tfrecord.py --csv_input=data/test_labels.csv  --output_path=test.record
"""
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 collections import namedtuple, OrderedDict

flags = tf.app.flags
flags.DEFINE_string('csv_input', '', 'Path to the CSV input')
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS


# TO-DO replace this with label map
def class_text_to_int(row_label):
    if row_label == 'raccoon':
        return 1
    else:
        None


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, path):
    with tf.gfile.GFile(os.path.join(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(class_text_to_int(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)
    path = os.path.join(os.getcwd(), 'images')
    examples = pd.read_csv(FLAGS.csv_input)
    grouped = split(examples, 'filename')
    for group in grouped:
        tf_example = create_tf_example(group, path)
        writer.write(tf_example.SerializeToString())

    writer.close()
    output_path = os.path.join(os.getcwd(), FLAGS.output_path)
    print('Successfully created the TFRecords: {}'.format(output_path))


if __name__ == '__main__':
    tf.app.run()

参考资料

[1] 对于谷歌开源的TensorFlow Object Detection API视频物体识别系统实现教程

[2] TensorFlow学习——Tensorflow Object Detection API（win10，CPU）

[3] How to train your own Object Detector with TensorFlow’s Object Detector API

[4] TensorFlow 之 物体检测（http://rensanning.iteye.com/category/374992）

---

1. See MSCOCO evaluation protocol.↩︎