Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)

结合上一章内容,本章节将结合实际需要,使用Tensorflow Object Detection API从头训练符合自己需求的图和表的检测分类模型.

需求说明:输入拍摄的文本页面图片,目标是训练一个可以正确检测图片上的图和表格并在图片画出其bounding box,之前用yolo也做过检测,不过效果不太好,会出现检测不准确甚至漏检情况,于是这次想尝试api的检测效果,模型使用ssd_mobilenet_v2.

1,训练数据准备:

因为之前用yolo做过检测,所以训练数据都放在一个文件夹下,另外还有一个.txt文件,.txt文件包含了图片路径,图标的bounding box及对应的类别,如图:

Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)_第1张图片

每行分别代表图片路径 xmin ymin xmax ymax label(0代表graph, 1代表table),一张图片可能有多个图表,所以每一行图片可能有多个xmin ymin xmax ymax label.我分别有两个这样的.txt文件,一个是训练集,一个是测试集,在后续生成tfrecode格式数据的时候分别生成即可.

2,tfrecord格式训练数据生成:

tfrecord训练数据的生成网上有很多,我参照官网代码生成自己的数据,tfrecord生成的时候可以选择生成一个或者多个.record文件,我参照最新代码生成多个文件,具体代码参考https://pan.baidu.com/s/1WO2OmMFyyBvqeOxtX9-xyQ,下载后将该代码放在下载的api的research目录下,然后python3 creratet_table_tf_record.py --data_dir=存放原始图片的路径 --output_dir=存放生成的数据的目录(训练和测试目录一样)执行即可.记得6行改成直接的图片路径,分别生成训练和测试数据时候114行table_test.record记得改(table_test.record, table_train.record),这是生成的训练和测试的名字.执行完后会在指定目录生成多个文件:

Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)_第2张图片

3,准备训练配置文件:

本次使用ssd_mobilenet_v2检测模型,所以将object_detection/samples/configs/ssd_mobilenet_v2_coco.config复制到存放刚刚生成的record文件的目录下,然后在该目录新建table_label_map.pbtxt文件,并添加内容:

Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)_第3张图片

因为tfrecord格式中的label都是从1开始,所以这从1开始(虽然txt中graph为0,table为1,但是在生成tfrecord的时候代码中已经把graph改成了1,table改成了2,所以和这里是一样的)

然后在该目录继续执行

sed -i "s|PATH_TO_BE_CONFIGURED|指定存放tfrecord数据的当前路径|g" ssd_mobilenet_v2_coco.config, (注意这是一行代码,例如
sed -i "s|PATH_TO_BE_CONFIGURED|/home/user/model/reaearch/objectdetect/data|g" ssd_mobilenet_v2_coco.config)

4,检测模型训练:

modify models/research/object_detection/model_main.py after imports add the following

tf.logging.set_verbosity(tf.logging.INFO)

在安装的api目录的research目录下执行(本人用的python3)

python3 object_detection/model_main.py \
    --pipeline_config_path=ssd_mobilenet_v2_coco.config文件路径 \
    --model_dir=存放训练模型的路径(没有会自己创建) \
    --num_train_steps=50000 \
    --num_eval_steps=2000 \
    --alsologtostderr

(num_train_steps是训练次数,本人是gpu训练,所以指定五万次)

5,tensorboard 查看训练过程:

执行tensorboard --logdir model_dir路径,然后打开tensorboard就能看到训练过程和训练图片,如图:

Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)_第4张图片

Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)_第5张图片

Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)_第6张图片

6,分类模型预测:

tensorboard object detect api中的models/research/objectdetect/目录下提供了一个object_detection_tutorial.ipynb文件可以做预测,本人对其进行改造,将其写入.py文件中进行预测,预测前先要将训练的模型转为.pb文件,上篇博客提到了如何操作,而且官网也有说明,很简单的,这里就不在阐述.此外,你也可以直接在.ipynb文件预测都可以的,预测结果:

Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)_第7张图片

Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)_第8张图片

 

Tensorflow Object Detection API 训练图表分类模型-ssd_mobilenet_v2(tfrecord数据准备+训练+测试)_第9张图片

可以看到预测不错,下一步将把该模型迁移到android平台进行测试....

你可能感兴趣的:(目标检测)