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

gloomyfish 图像算法工程师  2020-10-16 11:37:22
开发环境
软件版本信息:
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()

运行结果如下:

...全文
1606 点赞 收藏 回复
写回复
回复
切换为时间正序
请发表友善的回复…
发表回复

还没有回复,快来抢沙发~

相关推荐
发帖
英特尔边缘计算技术
创建于2007-08-27

450

社区成员

英特尔® 边缘计算,聚焦于边缘计算、AI、IoT等领域,为开发者提供丰富的开发资源、创新技术、解决方案与行业活动。
申请成为版主
帖子事件
创建了帖子
2020-10-16 11:37
社区公告
暂无公告