本文借助 Tensorflow Object detection API 开源框架和 MobileNet V2—SSD 算法，阐述如何创建、训练自定义车道线数据集，并获得 LDW 目标检测模型。
一、前期准备工作
前期准备工作主要包含模型下载、环境配置以及 py 文件生成等，详细内容如下：
1）模型下载 —— API 是基于 TensorFlow 构造的开源框架，易于构建、训练、部署目标检测模型，下载网址：https://github.com/tensorflow/models；模型下载后解压，并将其放置于指定位置，方便后期进行模型训练。其次，Protoc 用于将 Tensorflow object detection API 模型文件中的 .pro 文件编译成 python 文件，protoc-3.4.0-win32 版本，下载网址：https://github.com/protocolbuffers/protobuf/releases，文件下载后解压，并将 bin 文件夹的路径添加到 PC 环境变量；之后打开 cmd，输入 protoc，输入如下信息即表示添加环境变量成功，如图 1.1 所示。
图 1.1 protoc 配置成功示意图：
2）环境配置 —— 通用软件平台为 Win 10、JetBrains PyCharm Community Edition 2019.2 x64、 Anaconda 3( 64 bit )。数据集格式转换的主要配置为 Python 3.7、TensorFlow 1.13.1；模型训练、生成与测试的主要配置为 Python 3.5、TensorFlow 1.10.0、TensorFlow-gpu 1.10.0。
3）py 文件生成 —— 将解压后的 API 文件命名为 model，其子路径model\research\object_detection\protos 中有很多的 .proto 文件；Protoc 可以将上述 .proto 文件编译成 py 文件，即在 cmd 下， 一路 cd 到 …\models\research 文件夹下，输入：protoc ./object_detection/protos/*.proto --python_out=.
即可快速编译所有 proto格式文件；*.proto 相当于匹配所有的 .proto 文件（本次总结文件已将生成的 py 文件一并释放，后期可提供）。编译之后的文件如图 1.2 所示。
图 1.2 protoc 生成 py 文件示意图：
4）补充配置环境 —— 将 API 解压后的目录 research 及其子目录 slim 添加到环境变量，如图 1.3 所示。
图 1.3 补充环境配置示意图：
二、自定义数据集生成与格式转换
自定义数据集分为标注软件的安装、数据集的标注与生成以及格式转换。
1）标注软件的安装 —— 下载 labelImg，并创建 data 文件夹，如图 2.1 所示；data 路径下创建 predefined_classes.txt，其内容为：输入自定义的分类类别 —— 本次检查目标为车道线，即为 “lane”。如图 2.2 所示。
图 2.1 为 labelImg-maste 解压图：
图 2.2 为 lane 自定义标签 示意图：
2）数据集的标注与生成 —— 打开 labelImg, 按快捷键 CTRL+R, 修改默认的输出文件的保存路径，接着点击 open dir，添加图片所在的路径，如图 2.3 所示。点击 edit->create rectbox 进行标记。标记完后点击 save 即可，点击 next page 进行下一张图片标记。最终生成的数据文件如图 2.4 所示（图片与 XML 文件的名称与数量务必一一对应）。
图 2.3 为 labellmg 图像标注示意图：
图 2.4 为 labellmg 图像标注生成文件示意图：
补充说明：本次车道线数据集来源于上海某高架路段，采用视频截屏方式作为数据源，总数为 362 张，每张图片约有 8 ~~ 12 个标注对象。
3）数据格式转换
TensorFlow object detection API 对训练数据的格式要求是 tfrecord 格式，需要将 labelImg 生成的车道线数据格式以此转换为 xml à csv à tfrecord。详细步骤（请重点关注红色加深字体）如下：
第一步：将标记生成 xml 文件转换为 csv 文件：model\research\object_detection 目录下分布创建 xml_to_csv.py、 generate_tfrecord.py 以及 images 文件夹（创建 train_and_eval_images —— 放置原始车道线数据图片及其对应的 xml 文件）。
第二步：修改 xml_to_csv.py 文件 code（请重点关注红色加深字体）如下所示。
//************************************************************//
第 33、34 行 Code 更改为输出 csv 的路径和输出文件名：
33 train_df.to_csv('images/train.csv', index=None)
34 eval_df.to_csv('images/eval.csv', index=None)
第 38 行 Code 更改为原始数据路径（亦可用相对路径），参数更改为 xml 文件所在的文件夹路径修改：
path = 'D:\\WPI\\WorkSources\\pycharmExample\\Q4API\\model\\research\\object_detection\\images\\train_and_eval_images'
在 pycharm 平台运行上述代码并生成两个 csv 文件，即为 train.csv、 eval.csv。
第三步：修改 generate_tfrecord.py 文件 code 如下所示。
第 34、37 行更改标签及其对应的数字：
def class_text_to_int(row_label):
if row_label == 'lane': #此处更改为自定义检测对象，与 predefined_classes.txt 标签内容一致
return 1 #label 对应的数字标签，其值从 1 开始设置
else:
None #二分类情况下，此处为背景标签，设置为 None 即可
//************************************************************//
generate_tfrecord.py 在 pycharm 终端分别运行如下命令：
python generate_tfrecord.py --csv_input=images/train.csv --image_dir=images/train_and_eval_images --output_path=images/train.record #生成 train.record
python generate_tfrecord.py --csv_input=images/ eval.csv --image_dir=images/train_and_eval_images --output_path=images/ eval.record #生成 eval.record
生成效果如图 2.5 所示。
图 2.5 为 csv、record 格式生成示意图：
三、模型生成与测试
生成模块分为补充配置、模型训练、pb 生成以及模型测试等，详细内容如下：
1）补充配置 —— 配置一：在指定目录（ …\model\dataset ）下创建 data 文件夹，并在其路径下创建 pbtxt 文件，内容如图 3.1 所示。文件中 id 和 name 分别代表被测对象的数字标签 label 及其名称的字符串文本，若为多目标检测， 本文件需要另起一行，同时 id 数字依次增 1,相应的 name 也要更新（字符串不能重名）。
图 3.1 为 pbtxt 内容示意图:
配置二：dataset 文件夹下创建 fine_tune_model 文件夹，下载 fine-tune 模型，网址是 https://github.com/tensorflow/models/tree/master/research/slim，如图 3.2 所示；解压后复制三个 .ckpt 文件放到 fine_tune_model 文件夹下，如图 3.3 所示。
图 3.2 fine_tune_model 下载 示意图：
图 3.3 fine_tune_model 示意图：
配置三：配置管道文件，找到目录research\object_detection\samples\configs\ssdlite_mobilenet_v2_coco.config 文件，将此文件复制到dataset/data 文件夹下，并将前面生成的 train.record、 eval.record 放置于 dataset/data 路径下。修改 ssdlite_mobilenet_v2_coco.config 文件 code （请重点关注红色加深字体）如下所示：
//************************************************************//
第 9 行 Code 更改为检测种类数量：
num_classes: 1 #此处为被检查对象的数量，本次仅对车道线检测，故而设置为 1
第 158 行更改为 fine_tune_model 位置与同类名称
fine_tune_checkpoint: "D:/WPI/WorkSources/pycharmExample/Q4API/model/dataset/fine_tune_model/model.ckpt"
第 165 行更改为训练次数：
num_steps: 200000 #此处为训练次数设置，根据需要设定所需数值
第 178 行更改为训练 record 文件输入路径
input_path: "D:/WPI/WorkSources/pycharmExample/Q4API/model/dataset/data/train.record"
第 180 行将 label_map_path 参数改为 pbtxt 文件的路径（可以是相对路径或绝对路径）
label_map_path: "D:/WPI/WorkSources/pycharmExample/Q4API/model/dataset/data/label_map.pbtxt"
第 184 行为 patch 参数，其数值不能太大，否则会出现显卡内存溢出等问题，建议从小到大，逐步调参
num_examples: 27 #此处为 patch 参数，其数值不能太大，否则会出现显卡内存溢出等问题，建议从小到大，逐步调参
第 192 行为测试 record 文件输入路径
input_path: "D:/WPI/WorkSources/pycharmExample/Q4API/model/dataset/data/eval.record"
第 194 行为 label_map_path 参数改为 pbtxt 文件的路径
label_map_path: "D:/WPI/WorkSources/pycharmExample/Q4API/model/dataset/data/label_map.pbtxt"
//************************************************************//
2）模型训练 —— 首先，找到 …\model\research\object_detection 目录下 trainer.py 更改为 train_Lite.py（名称根据自己需要任意命名），修改上述文件的 code 如下：
//************************************************************//
第 76 行第二个参数为模型输出路径，可以根据需要定义相对路径或绝对路径
flags.DEFINE_string('train_dir', '..\\..\\dataset\\ssdlite_train_output_dir',
'Directory to save the checkpoints and training summaries.')
第 79 行第二个参数为模型训练所需的配合文件路径，可以定义相对路径或绝对路径
flags.DEFINE_string('pipeline_config_path', r'D:\WPI\WorkSources\pycharmExample\Q4API\model\dataset\data\ssdlite_mobilenet_v2_coco.config',
'Path to a pipeline_pb2.TrainEvalPipelineConfig config '
'file. If provided, other configs are ignored')
//************************************************************//
至此，运行脚本修改完毕，运行上述 py 文件进行训练（注意此处选择的配置），训练过程如图 3.4 所示。模型训练 20 万次之后如图 3.5 所示。
图 3.4 模型训练示意图：
图 3.5 模型训练完成示意图：
3）pb 生成 —— 将 …\model\research\object_detection 路径下的 export_inference_graph.py 复制到 ..\model\dataset 路径下，修改 code （请重点关注红色加深字体）如下所示：
//************************************************************//
第 118 行第二个参数为 pb 转换所需的模型配置文件路径（可以是绝对路径或相对路径）
flags.DEFINE_string('pipeline_config_path', 'D:\\WPI\\WorkSources\\pycharmExample\\Q4API\\model\\dataset\\data\\ssdlite_mobilenet_v2_coco.config',
'Path to a pipeline_pb2.TrainEvalPipelineConfig config '
'file.')
第 121 行第二个参数为模型训练后生成的依赖文件路径（可以是绝对路径或相对路径）
flags.DEFINE_string('trained_checkpoint_prefix', 'D:\\WPI\\WorkSources\\pycharmExample\\Q4API\\model\\dataset\\ssdlite_train_output_dir\\model.ckpt-200000',
'Path to trained checkpoint, typically of the form '
'path/to/model.ckpt')
第 124 行第二个参数为 pb 输出路径（可以是绝对路径或相对路径）
flags.DEFINE_string('output_directory', 'D:\\WPI\\WorkSources\\pycharmExample\\Q4API\\model\\dataset\\ssdlite_train_output_dir',
'Path to write outputs.')
//************************************************************//
Pb 文件生成效果如图 3.6 所示。
图 3.6 pb 文件生成效果图：
4）模型测试 —— 将生成的 pb 模型文件对视频中的车道线进行检测，object_detection_video.py 为检测运行脚本，修改内容（请重点关注红色加深字体）如下：
//************************************************************//
第 35 行修改成车道线 pb 模型路径（可以是绝对路径或相对路径）
PATH_TO_CKPT = 'D:\\WPI\\WorkSources\\pycharmExample\\Q4API\\model\\dataset\\ssdlite_train_output_dir\\frozen_inference_graph.pb'
第 40 行修改成自己的 pbtxt 文件（可以是绝对路径或相对路径）
PATH_TO_LABELS = os.path.join('D:\WPI\WorkSources\pycharmExample\Q4API\model\dataset\data', 'label_map.pbtxt')
第43 行其数值参数为被检查对象的数量，本次检查只有车道线，故而数值是 1
NUM_CLASSES = 1
第 65 行更改为自己的视频，其格式可以为 avi、mp4
capture = cv2.VideoCapture('./test_video/lane_17_21.avi')
//************************************************************//
本次车道线检测效果如图 3.7 所示。
图 3.7 车道线检测效果图：
四、总结
本文完成 Tensorflow Object detection API 开源框架配合 MobileNet_v2 – SSD 算法实现车道线的目标检测，其检测精度可达 95% 以上；鉴于数据集来自上海市某高架桥，其鲁棒性还有待提高，后期 WPI ATU 将对模型进行修改、裁剪、调参、提高帧率等优化操作。欢迎在大大通上关注我，给我留言。
参考资料
- 途径北海道 —— Tensorflow object detection API训练自己的目标检测模型
链接：https://blog.csdn.net/zj1131190425/article/details/80711857
评论