发布于2019-08-19 10:26 阅读(1155) 评论(0) 点赞(0) 收藏(4)
opencv中的dnn模块很强大,可以加载现在主流的深度学习模型,darknet、caffe、tensorflow等,我本人最近在使用tensorflow进行目标检测,用的是ssdlite-mobilenet-v2。tensorflow的object_detection api使用起来很方便,而且大佬研究出来的模型也都没毛病,效果很好,换换数据集,调调参数就ok。训练过程不再本篇博客的讨论范围内,假定现在已经训练完毕,并且生成了pb文件。
一般我们恢复模型都是要用到tensorflow,大致代码如下所示:
- import numpy as np
- import tensorflow as tf
- import cv2 as cv
-
- # Read the graph.
- with tf.gfile.FastGFile('frozen_inference_graph.pb', 'rb') as f:
- graph_def = tf.GraphDef()
- graph_def.ParseFromString(f.read())
-
- with tf.Session() as sess:
- # Restore session
- sess.graph.as_default()
- tf.import_graph_def(graph_def, name='')
-
- # Read and preprocess an image.
- img = cv.imread('example.jpg')
- rows = img.shape[0]
- cols = img.shape[1]
- inp = cv.resize(img, (300, 300))
- inp = inp[:, :, [2, 1, 0]] # BGR2RGB
-
- # Run the model
- out = sess.run([sess.graph.get_tensor_by_name('num_detections:0'),
- sess.graph.get_tensor_by_name('detection_scores:0'),
- sess.graph.get_tensor_by_name('detection_boxes:0'),
- sess.graph.get_tensor_by_name('detection_classes:0')],
- feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)})
-
- # Visualize detected bounding boxes.
- num_detections = int(out[0][0])
- for i in range(num_detections):
- classId = int(out[3][0][i])
- score = float(out[1][0][i])
- bbox = [float(v) for v in out[2][0][i]]
- if score > 0.3:
- x = bbox[1] * cols
- y = bbox[0] * rows
- right = bbox[3] * cols
- bottom = bbox[2] * rows
- cv.rectangle(img, (int(x), int(y)), (int(right), int(bottom)), (125, 255, 51), thickness=2)
-
- cv.imshow('TensorFlow MobileNet-SSD', img)
- cv.waitKey()
如果电脑没有tensorflow想要恢复模型怎么办?那就可以考虑下opencv了,opencv加载tensorflow模型需要pb文件和pbtxt文件,pbtxt是可以根据pb文件生成的,在opencv的源代码中进入sample/dnn 文件夹中,由于我们的模型是ssd因此找到以下两个文件
在命令行执行
ssdlite.pb是自己训练的pb模型,pipeline是训练时的配置文件,生成pb时会有,ssdlite.pbtxt就是生成的文件用于opencv恢复模型
但是我本人在这里运行时报错了,如下图所示:
在github中有人提到训练完的模型需要转换一下
- import tensorflow as tf
- from tensorflow.tools.graph_transforms import TransformGraph
-
- with tf.gfile.FastGFile('ssdlite.pb', 'rb') as f:
- graph_def = tf.GraphDef()
- graph_def.ParseFromString(f.read())
- graph_def = TransformGraph(graph_def, ['image_tensor'], ['detection_boxes', 'detection_classes', 'detection_scores', 'num_detections'], ['sort_by_execution_order'])
- with tf.gfile.FastGFile('ssdlite_new.pb', 'wb') as f:
- f.write(graph_def.SerializeToString())#保存新的模型
转化完pb模型之后再运行上面的代码,成功生成pbtxt文件
然后用opencv恢复模型
大致代码如下:
- import cv2 as cv
-
- cvNet = cv.dnn.readNetFromTensorflow('frozen_inference_graph.pb', 'graph.pbtxt')
-
- img = cv.imread('test.jpg')
- rows = img.shape[0]
- cols = img.shape[1]
- cvNet.setInput(cv.dnn.blobFromImage(img, size=(300, 300), swapRB=True, crop=False))
- cvOut = cvNet.forward()
-
- for detection in cvOut[0,0,:,:]:
- score = float(detection[2])
- if score > 0.3:
- left = detection[3] * cols
- top = detection[4] * rows
- right = detection[5] * cols
- bottom = detection[6] * rows
- cv.rectangle(img, (int(left), int(top)), (int(right), int(bottom)), (23, 230, 210), thickness=2)
-
- cv.imshow('img', img)
- cv.waitKey()
在此需要感谢dkurt大神的转换代码,顺便提一下ssdlite的速度是真的快,pc端cpu下大约60ms一帧,如果需要移植到arm平台上,可借助腾讯的ncnn,在移动端加速效果也是很棒的。
参考的一些资料:
https://github.com/opencv/opencv/wiki/TensorFlow-Object-Detection-API
作者:爸爸去挣钱我去幼儿园
链接:https://www.pythonheidong.com/blog/article/48827/6e671ad30b2f7d55fc2f/
来源:python黑洞网
任何形式的转载都请注明出处,如有侵权 一经发现 必将追究其法律责任
昵称:
评论内容:(最多支持255个字符)
---无人问津也好,技不如人也罢,你都要试着安静下来,去做自己该做的事,而不是让内心的烦躁、焦虑,坏掉你本来就不多的热情和定力
Copyright © 2018-2021 python黑洞网 All Rights Reserved 版权所有,并保留所有权利。 京ICP备18063182号-1
投诉与举报,广告合作请联系vgs_info@163.com或QQ3083709327
免责声明:网站文章均由用户上传,仅供读者学习交流使用,禁止用做商业用途。若文章涉及色情,反动,侵权等违法信息,请向我们举报,一经核实我们会立即删除!