社区首页 >专栏 >Tensorflow + OpenCV4 安全帽检测模型训练与推理

Tensorflow + OpenCV4 安全帽检测模型训练与推理

OpenCV学堂

发布于 2020-09-08 09:03:27

2.5K00

代码可运行

文章被收录于专栏：贾志刚-OpenCV学堂贾志刚-OpenCV学堂

运行总次数：0

代码可运行

开发环境 ·

软件版本信息：

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

如何导出PB模型到OpenCV DNN支持看这里：

干货 | tensorflow模型导出与OpenCV DNN中使用

使用OpenCV DNN调用模型

在OpenCV DNN中直接调用训练出来的模型完成自定义对象检测，这里需要特别说明一下的，因为在训练阶段我们选择了模型支持600~1024保持比率的图像输入。所以在推理预测阶段，我们可以直接使用输入图像的真实大小，模型的输出格式依然是1x1xNx7，按照格式解析即可得到预测框与对应的类别。最终的代码实现如下：

import cv2 as cv

labels = ['hat', 'person']
model = "D:/safehat_train/models/train/frozen_inference_graph.pb"
config = "D:/safehat_train/models/train/frozen_inference_graph.pbtxt"

# 读取测试图像
image = cv.imread("D:/123.jpg")
h, w = image.shape[:2]
cv.imshow("input", image)

# 加载模型，执行推理
net = cv.dnn.readNetFromTensorflow(model, config)
blob = cv.dnn.blobFromImage(cv.resize(image, (w, h)), swapRB=True, crop=False)
net.setInput(blob)
detectOut = net.forward()

# 解析输出
classIds = []
confidences = []
boxes = []
for detection in detectOut[0,0,:,:]:
    score = detection[2]
    if score > 0.4:
        left = detection[3]*w
        top = detection[4]*h
        right = detection[5]*w
        bottom = detection[6]*h
        classId = int(detection[1]) + 1
        classIds.append(classId)
        boxes.append([int(left), int(top), int(right), int(bottom)])
        confidences.append(float(score))

# 非最大抑制
nms_indices = cv.dnn.NMSBoxes(boxes, confidences, 0.4, 0.4)
for i in range(len(nms_indices)):
    index = nms_indices[i][0]
    box = boxes[index]
    cid = classIds[index]
    if cid == 1:
        cv.rectangle(image, (box[0], box[1]), (box[2], box[3]), (140, 199, 0), 4, 8, 0)
    else:
        cv.rectangle(image, (box[0], box[1]), (box[2], box[3]), (255, 0, 255), 4, 8, 0)
    cv.putText(image, labels[cid-1], (box[0], box[1]), cv.FONT_HERSHEY_SIMPLEX, 0.75, (255, 0, 0), 2)

# 显示输出
cv.imshow("safetyhat-detection-demo", image)
cv.imwrite("D:/result123.png", image)
cv.waitKey(0)
cv.destroyAllWindows()