5.2 TensorFlow Object Detection API
2017年6月,Google公司开放了TensorFlow Object Detection API。这个项目使用TensorFlow实现了大多数深度学习目标检测框架,其中就包括Faster R-CNN。在本节中,首先介绍如何安装TensorFlow Object Detection API,再介绍如何使用已经训练好的模型进行物体检测,最后介绍如何训练自己的模型。
5.2.1 安装TensorFlow Object Detection API
在GitHub上,TensorFlow Object Detection API是存放在tensorflow/models项目(地址:https://github.com/tensorflow/models)下的。可以通过git来下载tensorflow/models:
git clone https://github.com/tensorflow/models.git。
下载tensorflow/models代码后,应该得到一个models文件夹。models文件夹中还有一个research文件夹。下面的安装命令都是以research文件夹为根目录执行的,所说的目录也都是以research文件夹为相对目录。
安装TensorFlow Object Detection API的步骤如下(以research文件夹为相对目录):
1.安装或升级protoc
在object_detection/protos/中,可以看到一些proto文件,需要使用protoc程序将这些proto文件编译为Python文件。TensorFlow Object Detection API必须使用2.6.0以上的protoc进行编译,否则会报错。可以使用命令protoc--version查看protoc的版本。如果发现版本低于2.6.0或运行命令错误,就需要安装或升级protoc。
安装或升级的方法是登录protobuf的发布页面:https://github.com/google/protobuf/releases下载已经编译好的文件包。
如图5-9所示,默认有多种版本的预编译包。需要自行找到对应机器应该使用的文件。如64位的ubuntu系统应下载文件protoc-3.3.0-linux-x86_64.zip,64位的OS X系统应下载文件protoc-3.3.0-osx-x86_64.zip。下载文件后解压,会得到一个protoc文件,将它复制到系统的可执行目录即可,如在ubuntu系统中,可以执行以下命令:
图5-9 protoc的源码和各个系统的预编译包
sudo cp bin/protoc /usr/bin/protoc
2.编译proto文件
使用protoc对proto文件进行编译。具体来说,应当在research文件夹下,运行下面的命令:
# From models/research protoc object_detection/protos/*.proto--python_out=.
运行完成后,可以检查object_detection/protos/文件夹,如果每个proto文件都生成了对应的以py为后缀的python源代码,就说明编译成功了。
3.将Slim加入PYTHONPATH
TensorFlow Object Detection API是以Slim为基础实现的,需要将Slim的目录加入PYTHONPATH后才能正确运行。具体来说,还是在research文件夹下,执行下面的命令:
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
执行命令完成后,可以使用python命令打开一个python shell,如果运行import slim成功则说明已经正确设置好了。
4.安装完成测试
在research文件夹下,执行:
python object_detection/builders/model_builder_test.py
这条命令会自动检查TensorFlow Object Detection API是否正确安装,如果出现下面的信息,说明已安装成功:
....... ------------------------------------------------------------------ Ran 7 tests in 0.019s OK
5.2.2 执行已经训练好的模型
TensorFlow Object Detection API默认提供了5个预训练模型,它们都是使用COCO数据集训练完成的,结构分别为SSD+MobileNet、SSD+Inception、R-FCN+ResNet101、Faster RCNN+ResNet101、Faster RCNN+Inception_ResNet。
如何使用这些预训练模型呢?官方已经给了一个用Jupyter Notebook编写好的例子。首先在research文件夹下,运行命令:jupyter-notebook。如果提示不存在该命令,可能是因为没有安装Jupyter Notebook,需要读者自行安装。
运行命令jupyter-notebook后,打开http://localhost:8888,接着打开object_detection文件夹,并单击object_detection_tutorial.ipynb运行示例文件,如图5-10所示。
图5-10 官方提供的示例object_detection_tutorial.ipynb
使用组合键“Shift+Enter”可以依次执行这些命令。在这里介绍Notebook中的命令,并给出相应的中文注释。
首先是导入一些需要的包和设置环境:
# 导入基本的包 import numpy as np import os import six.moves.urllib as urllib import sys import tarfile import tensorflow as tf import zipfile from collections import defaultdict from io import StringIO from matplotlib import pyplot as plt from PIL import Image # 这条命令让在使用matplotlib绘图时,不再使用窗口展示出来,而是直接在notebook中显 示%matplotlib inline # 将上层目录导入进来,这样才可以执行这下面的两条导入命令 sys.path.append("..") from utils import label_map_util from utils import visualization_utils as vis_util
导入包后,设置需要使用的模型:
# 使用的模型名称,下面会下载这个模型 MODEL_NAME='ssd_mobilenet_v1_coco_11_06_2017' MODEL_FILE=MODEL_NAME+'.tar.gz' DOWNLOAD_BASE='http://download.tensorflow.org/models/object_detection/' # frozen_inference_graph.pb文件就是在后面需要导入的文件,它保存了网络的结构和数据 PATH_TO_CKPT=MODEL_NAME+'/frozen_inference_graph.pb' # mscoco_label_map.pbtxt文件中保存了index到类别名的映射 # 如神经网络的预测类别是5,必须要通过这个文件才能知道index具体对应的类别是什么 # mscoco_label_map.pbtxt文件就保存在object_detection/data文件夹下,读者可以自 行打开查阅 PATH_TO_LABELS=os.path.join('data', 'mscoco_label_map.pbtxt') NUM_CLASSES=90
接下来下载预训练模型,根据网络环境的不同,下载的时间可能会有长有短。
opener=urllib.request.URLopener() opener.retrieve(DOWNLOAD_BASE+MODEL_FILE, MODEL_FILE) tar_file=tarfile.open(MODEL_FILE) for file in tar_file.getmembers(): file_name=os.path.basename(file.name) if 'frozen_inference_graph.pb' in file_name: tar_file.extract(file, os.getcwd())
这里程序组合了DOWNLOAD_BASE和MODEL_FILE两个变量得到了下载地址。很显然,由于MODEL_FILE的值是‘ssd_mobilenet_v1_coco_11_06_2017',因此下载的模型为SSD+MobileNet。如何下载其他预训练模型并执行呢?实际上也很简单,可以打开TensorFlow detection model zoo,找到其他模型的下载地址。根据这些下载地址,只需要改变MODEL_FILE的值就可以下载不同的模型。为方便读者查阅,在此处也列出这些值:
MODEL_NAME='ssd_inception_v2_coco_11_06_2017' MODEL_NAME='rfcn_resnet101_coco_11_06_2017' MODEL_NAME='faster_rcnn_resnet101_coco_11_06_2017' MODEL_NAME='faster_rcnn_inception_resnet_v2_atrous_coco_11_06_2017'
回到示例代码,下载模型后,程序就直接将它读取到默认的计算图中(实际读取的是frozen_inference_graph.pb文件),使用的代码如下所示:
# 新建一个图 detection_graph=tf.Graph() with detection_graph.as_default(): od_graph_def=tf.GraphDef() # PATH_TO_CKPT指向了文件frozen_inference_graph.pb 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='')
在进行真正的检测之前,还得定义一些帮助函数:
# 这部分代码的功能是将神经网络检测得到的index(数字)转换为类别名(字符串) label_map=label_map_util.load_labelmap(PATH_TO_LABELS) categories=label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True) category_index=label_map_util.create_category_index(categories) # 这个函数也是一个方便使用的帮助函数,功能是将图片转换为Numpy数组的形式 def load_image_into_numpy_array(image): (im_width, im_height)=image.size return np.array(image.getdata()).reshape( (im_height, im_width, 3)).astype(np.uint8)
下面开始检测图片!先定义要检测的图片:
# 只检测了两张图片 PATH_TO_TEST_IMAGES_DIR='test_images' TEST_IMAGE_PATHS=[ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'. format(i)) for i in range(1, 3) ] # 输出图像的大小(单位是in) IMAGE_SIZE=(12, 8)
TEST_IMAGE_PATHS是一个列表,它保存了需要检测的图片。这里检测的图片是官方提供的示例图片。如果想要检测自己的图片,只要将这些图片的路径以列表形式保存在TEST_IMAGE_PATHS中就可以了。
最后是检测的代码,同样给出中文注释:
# detection_graph是之前定义好的计算图,已经将模型导入到内存中了,此处直接使用 with detection_graph.as_default(): with tf.Session(graph=detection_graph) as sess: for image_path in TEST_IMAGE_PATHS: image=Image.open(image_path) # 将图片转换为Numpy的形式 image_np=load_image_into_numpy_array(image) # 将图片扩展一维,最后进入神经网络的图片的格式应该为[1, ? , ? , 3] image_np_expanded=np.expand_dims(image_np, axis=0) image_tensor=detection_graph.get_tensor_by_name('image_tensor:0') # boxes变量存放了所有检测框 boxes=detection_graph.get_tensor_by_name('detection_boxes:0') # score表示每个检测结果的confidence scores=detection_graph.get_tensor_by_name('detection_scores:0') # classes表示每个框对应的类别 classes=detection_graph.get_tensor_by_name('detection_classes:0') # num_detections表示检测框的个数 num_detections=detection_graph.get_tensor_by_name('num_detections: 0') # 使用sess.run,真正开始计算 (boxes, scores, classes, num_detections)=sess.run( [boxes, scores, classes, num_detections], feed_dict={image_tensor: image_np_expanded}) # 对得到的检测结果进行可视化 vis_util.visualize_boxes_and_labels_on_image_array( image_np, np.squeeze(boxes), np.squeeze(classes).astype(np.int32), np.squeeze(scores), category_index, use_normalized_coordinates=True, line_thickness=8) plt.figure(figsize=IMAGE_SIZE) plt.imshow(image_np)
两张示例图片的检测效果如图5-11所示。
图5-11 示例图片的检测效果
5.2.3 训练新的模型
以VOC 2012数据集为例,介绍如何使用TensorFlow Object Detection API训练新的模型。VOC 2012是VOC 2007数据集的升级版,一共有11530张图片,每张图片都有标注,标注的物体包括人、动物(如猫、狗、鸟等)、交通工具(如车、船飞机等)、家具(如椅子、桌子、沙发等)在内的20个类别。图5-12展示了VOC 2012中的一张图片。
图5-12 VOC 2012数据集的示例图片
首先下载数据集,并将其转换为tfrecord格式。VOC 2012数据集的下载地址为http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar。
为了不影响代码的结构,不妨在object_detection文件夹中,再新建一个voc文件夹,并将下载的数据集压缩包复制至voc/中。解压后,就得到一个VOCdevkit文件夹,最终的文件夹结构应该为
research/ object_detection/ voc/ VOCdevkit/ VOC2012/ JPEGImages/ 2007_000027.jpg 2007_000032.jpg 2007_000033.jpg 2007_000039.jpg 2007_000042.jpg ……………… Annotations/ 2007_000027.xml 2007_000032.xml 2007_000033.xml 2007_000039.xml 2007_000042.xml ……………… ………………
JPEGImages文件中存储了所有的图像数据。对于每一张图片,都在Annotations文件夹中有其物体框的标注。
在object_detection文件夹中,执行以下命令可以将VOC 2012数据集转换为tfrecord格式,转换好的tfrecord保存在新建的voc文件夹下,分别为pascal_train.record和pascal_val.record:
python create_pascal_tf_record.py--data_dir voc/VOCdevkit/--year=VOC2012 --set=train--output_path=voc/pascal_train.record python create_pascal_tf_record.py--data_dir voc/VOCdevkit/--year=VOC2012 --set=val--output_path=voc/pascal_val.record
此外,将pascal_label_map.pbtxt数据复制到voc文件夹下:
cp data/pascal_label_map.pbtxt voc/
这里的转换代码是为VOC 2012数据集提前编写好的。如果读者希望使用自己的数据集,有两种方法,第一种方法是修改自己的数据集的标注格式,使其和VOC 2012一模一样,然后就可以直接使用create_pascal_tf_record.py脚本转换了,另外一种方法是修改create_pascal_tf_record.py,对读取标签的代码进行修改。
回到VOC 2012数据集的训练。下载数据集后,需要选择合适的模型。这里以Faster R-CNN+Inception_ResNet_v2模型为例进行介绍。首先下载在COCO上预训练的Faster R-CNN+Inception_ResNet_v2模型。解压后得到frozen_inference_graph.pb、graph.pbtxt、model.ckpt.data-00000-of-00001、model.ckpt.index、model.ckpt.meta 5个文件。在voc文件夹中新建一个pretrained文件夹,并将这5个文件复制进去。
TensorFlow Object Detection API是依赖一个特殊的设置文件进行训练的。在object_detection/samples/configs/文件夹下,有一些设置文件的示例。可以参考faster_rcnn_inception_resnet_v2_atrous_pets.config文件创建的设置文件。先将faster_rcnn_inception_resnet_v2_atrous_pets.config复制一份到voc文件夹下:
cp samples/configs/faster_rcnn_inception_resnet_v2_atrous_pets.config \ voc/voc.config
voc.config一共有7处需要修改的地方:
· 第一处为num_classes,需要将它修改为VOC 2012中的物体类别数,即20类。
· 第二处为eval_config中的num_examples,它表示在验证阶段需要执行的图片数量,修改为VOC 2012验证集的图片数5823(可以在create_pascal_tf_record.py中,输出对应的examples_list的长度,就可以知道这个大小)。
· 还有5处为所有含有PATH_TO_BE_CONFIGURED的地方。这些地方需要修改为自己的目录。它们应该分别被修改为
fine_tune_checkpoint: "voc/pretrained/model.ckpt" train_input_reader: { tf_record_input_reader { input_path: "voc/pascal_train.record" } label_map_path: "voc/pascal_label_map.pbtxt" } eval_input_reader: { tf_record_input_reader { input_path: "voc/pascal_val.record" } label_map_path: "voc/pascal_label_map.pbtxt" shuffle: false num_readers: 1 }
最后,在voc文件夹中新建一个train_dir作为保存模型和日志的目录,使用下面的命令就可以开始训练了:
python train.py--train_dir voc/train_dir/--pipeline_config_path voc/voc.config
训练的日志和最终的模型都会被保存在train_dir中,因此,同样可以使用TensorBoard来监控训练情况:
tensorboard--logdir voc/train_dir/
需要注意的是,如果发生内存和显存不足报错的情况,除了换用较小的模型进行训练外,还可以修改配置文件中的以下部分:
image_resizer { keep_aspect_ratio_resizer { min_dimension: 600 max_dimension: 1024 } }
这个部分表示将输入图像进行等比例缩放再开始训练,缩放后最大边长为1024,最小边长为600。可以将这两个数值改小(如分别改成512和300),使用的显存就会变小。不过这样做也很有可能导致模型的精度下降,读者还需根据自己的情况选择适合的处理方法。
5.2.4 导出模型并预测单张图片
如何将train_dir中的checkpoint文件导出并用于单张图片的目标检测?TensorFlow Object Detection API提供了一个export_inference_graph.py脚本用于导出训练好的模型。具体方法是执行:
python export_inference_graph.py \ --input_type image_tensor \ --pipeline_config_path voc/voc.config \ --trained_checkpoint_prefix voc/train_dir/model.ckpt-1582 --output_directory voc/export/
其中,model.ckpt-1582表示使用第1582步保存的模型。读者需要根据voc/train_dir/里实际保存的checkpoint,将1582改为合适的数值。导出的模型是voc/export/ frozen_inference_graph.pb文件。
读者可以参考第5.2.2节中Jupyter Notebook的代码,自行编写利用导出模型对单张图片做目标检测的脚本。首先去掉无用的下载模型的部分,然后将PATH_TO_CKPT的值赋值为“voc/export/frozen_inference_graph.pb”,即导出的模型文件。将PATH_TO_LABELS修改为“voc/pascal_label_map.pbtxt”,即各个类别的名称。其他代码都可以不改变,识别的效果如图5-13所示。
图5-13 训练深度学习目标检测模型(详见彩插)