【干货】tensorflow + OpenVINO实现安全帽检测

[gloomyfish]

开发环境
软件版本信息：

Windows10 64位
Tensorflow1.15
Tensorflow object detection API 1.x
Python3.6.5
VS2015 VC++
CUDA10.0

硬件：

CPUi7
GPU 1050ti

如何安装tensorflow object detection API框架，看这里：
Tensorflow Object Detection API 终于支持tensorflow1.x与tensorflow2.x了
数据集处理与生成
首先需要下载数据集，下载地址为：

https://pan.baidu.com/s/1UbFkGm4EppdAU660Vu7SdQ

总计7581张图像，基于Pascal VOC2012完成标注。分为两个类别，分别是安全帽与人（hat与person），json格式如下：

item {
  id: 1
  name: 'hat'
}

item {
  id: 2
  name: 'person'
}

数据集下载之后，并不能被tensorflow object detection API框架中的脚本转换为tfrecord，主要是有几个XML跟JPEG图像格式错误，本人经过一番磨难之后把它们全部修正了。修正之后的数据运行下面两个脚本即可生成训练集与验证集的tfrecord数据，命令行如下：

这里需要注意的是create_pascal_tf_record.py 脚本的165行把
'aeroplane_' + FLAGS.set + '.txt')
修改为：
FLAGS.set + '.txt')
原因是这里的数据集没有做分类train/val。所以需要修改一下，修改完成之后保存。运行上述的命令行，就可以正确生成tfrecord，否则会遇到错误。

模型训练
基于faster_rcnn_inception_v2_coco对象检测模型实现迁移学习，首先需要配置迁移学习的config文件，对应的配置文件可以从:
research\object_detection\samples\configs
中发现，发现文件:
faster_rcnn_inception_v2_coco.config
之后，修改配置文件的中相关部分，关于如何修改，修改什么，可以看这里：


修完完成之后，在D盘下新建好几个目录之后，执行下面的命令行参数：

就会开始训练，总计训练40000 step。训练过程中可以通过tensorboard查看训练结果：
模型导出
完成了40000 step训练之后，就可以看到对应的检查点文件，借助tensorflow object detection API框架提供的模型导出脚本，可以把检查点文件导出为冻结图格式的PB文件。相关的命令行参数如下：

得到pb文件之后，使用OpenCV4.x中的tf_text_graph_faster_rcnn.py脚本，转换生成graph.pbtxt配置文件。最终得到：

- frozen_inference_graph.pb
- frozen_inference_graph.pbtxt

转为OpenVINO IR文件

python mo_tf.py 
--input_model D:\safehat_train\models\train\frozen_inference_graph.pb /
--transformations_config extensions\front\tf\faster_rcnn_support_api_v1.15.json /
--tensorflow_object_detection_api_pipeline_config D:\safehat_train\models\pipeline.config /
--input_shape [1,600,600,3] /
--reverse_input_channels

使用模型实现，安全帽检测

from __future__ import print_function
import cv2
import time
import logging as log
from openvino.inference_engine import IECore
labels = ["hat", "person"]


def safehat_detection():
    model_xml = "D:/projects/opencv_tutorial/data/models/safetyhat/frozen_inference_graph.xml"
    model_bin = "D:/projects/opencv_tutorial/data/models/safetyhat/frozen_inference_graph.bin"

    log.info("Creating Inference Engine")
    ie = IECore()
    # Read IR
    net = ie.read_network(model=model_xml, weights=model_bin)

    log.info("Preparing input blobs")
    input_it = iter(net.input_info)
    input1_blob = next(input_it)
    input2_blob = next(input_it)
    print(input2_blob)
    out_blob = next(iter(net.outputs))

    # Read and pre-process input images
    print(net.input_info[input1_blob].input_data.shape)
    print(net.input_info[input2_blob].input_data.shape)

    image = cv2.imread("D:/safehat/test/1.jpg")
    ih, iw, ic = image.shape
    image_blob = cv2.dnn.blobFromImage(image, 1.0, (600, 600), (0, 0, 0), False)

    # Loading model to the plugin
    exec_net = ie.load_network(network=net, device_name="CPU")

    # Start sync inference
    log.info("Starting inference in synchronous mode")
    inf_start1 = time.time()
    res = exec_net.infer(inputs={input1_blob:[[600, 600, 1.0]], input2_blob: [image_blob]})
    inf_end1 = time.time() - inf_start1
    print("inference time(ms) : %.3f" % (inf_end1 * 1000))

    # Processing output blob
    res = res[out_blob]
    for obj in res[0][0]:
        if obj[2] > 0.5:
            index = int(obj[1])
            xmin = int(obj[3] * iw)
            ymin = int(obj[4] * ih)
            xmax = int(obj[5] * iw)
            ymax = int(obj[6] * ih)
            cv2.putText(image,labels[index-1], (xmin, ymin), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 0), 2)
            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2, 8, 0)
    cv2.imshow("safetyhat detection", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    safehat_detection()

运行结果如下：