干貨 | tensorflow模型導(dǎo)出與OpenCV DNN中使用
點(diǎn)擊上方“小白學(xué)視覺”，選擇加"星標(biāo)"或“置頂”
重磅干貨，第一時(shí)間送達(dá)
本文轉(zhuǎn)自｜OpenCV學(xué)堂
OpenCV DNN模塊
Deep Neural Network - DNN 是OpenCV中的深度神經(jīng)網(wǎng)絡(luò)模塊，支持基于深度學(xué)習(xí)模塊前饋網(wǎng)絡(luò)運(yùn)行、實(shí)現(xiàn)圖像與視頻場景中的
圖像分類
對象檢測
圖像分割
其模型導(dǎo)入與加載的相關(guān)API支持以下深度學(xué)習(xí)框架
tensorflow - readNetFromTensorflow
caffe - readNetFromCaffe
pytorch - readNetFromTorch
darknet - readNetFromDarknet
OpenCV3.4.1以上版本支持tensorflow1.11版本以上的對象檢測框架(object detetion)模型導(dǎo)出使用，當(dāng)前支持的模型包括以下：
也就是說通過tensorflow object detection API框架進(jìn)行遷移學(xué)習(xí)訓(xùn)練模型，導(dǎo)出預(yù)測圖之后，可以通過OpenCV3.4.1以上版本提供幾個(gè)python腳本導(dǎo)出graph配置文件，然后就可以在OpenCV DNN模塊中使用tensorflow相關(guān)的模型了。感覺十分方便，下面就按照操作走一波！
使用tensorflow模型
根據(jù)tensorflow中遷移學(xué)習(xí)或者下載預(yù)訓(xùn)練模型不同，OpenCV DNN 模塊提供如下可以使用腳本生成對應(yīng)的模型配置文件
tf_text_graph_ssd.py
tf_text_graph_faster_rcnn.py
tf_text_graph_mask_rcnn.py
這是因?yàn)?，OpenCV DNN需要根據(jù)text版本的模型描述文件來解析tensorflow的pb文件，實(shí)現(xiàn)網(wǎng)絡(luò)模型加載。對SSD對象檢測模型，生成模型描述文件運(yùn)行以下命令行即可(在一行執(zhí)行)：
python tf_text_graph_ssd.py
--input /path/to/model.pb
--config /path/to/example.config
--output /path/to/graph.pbtxt
以MobileNet-SSD v2版本為例，首先下載該模型，解壓縮以后會(huì)發(fā)現(xiàn)里面有一個(gè)frozen_inference_graph.pb文件，使用tensorflow加載預(yù)測圖進(jìn)行預(yù)測的代碼如下：
import tensorflow as tf
import cv2 as cv
# Read the graph.
model_dir = 'D:/tensorflow/ssd_mobilenet_v2_coco_2018_03_29/frozen_inference_graph.pb'
with tf.gfile.FastGFile(model_dir, '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('D:/images/objects.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()
運(yùn)行結(jié)果如下：
基于frozen_inference_graph.pb生成graph.pbtxt模型配置文件，命令行運(yùn)行截圖如下：
使用OpenCV DNN模塊加載tensorflow模型(frozen_inference_graph.pb與graph.pbtxt)，實(shí)現(xiàn)預(yù)測圖使用的代碼如下(注意此時(shí)不需要依賴tensorflow)：
import cv2 as cv
model_path = 'D:/tensorflow/ssd_mobilenet_v2_coco_2018_03_29/frozen_inference_graph.pb'
config_path = 'D:/tensorflow/ssd_mobilenet_v2_coco_2018_03_29/graph.pbtxt'
net = cv.dnn.readNetFromTensorflow(model_path, config_path)
frame = cv.imread('D:/images/objects.jpg')
rows = frame.shape[0]
cols = frame.shape[1]
net.setInput(cv.dnn.blobFromImage(frame, size=(300, 300), swapRB=True, crop=False))
cvOut = net.forward()
print(cvOut)
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(frame, (int(left), int(top)), (int(right), int(bottom)), (23, 230, 210), thickness=2)
cv.imshow('opencv-dnn-ssd-detect', frame)
cv.waitKey()
運(yùn)行結(jié)果如下(跟tensorflow中的運(yùn)行結(jié)果完全一致，OpenCV DNN果然靠譜)：
OpenCV DNN 行人檢測
本人嘗試了基于tensorflow object detection API使用MobileNet-SSD v2遷移學(xué)習(xí)實(shí)現(xiàn)自定義數(shù)據(jù)集訓(xùn)練，導(dǎo)出預(yù)測圖之后，使用OpenCV DNN模塊的python腳本生成對象的圖配置文件graph.pbtxt，通過OpenCV加載模型使用，實(shí)時(shí)預(yù)測，最后上一張運(yùn)行結(jié)果圖：
OpenCV DNN調(diào)用代碼如下
import cv2 as cv
inference_pb = "D:/pedestrian_data/export_pb/frozen_inference_graph.pb";
graph_text = "D:/pedestrian_data/export_pb/graph.pbtxt";
# load tensorflow model
net = cv.dnn.readNetFromTensorflow(inference_pb, graph_text)
image = cv.imread("D:/python/Pedestrian-Detection/test_images/3600.jpg")
h = image.shape[0]
w = image.shape[1]
# 獲得所有層名稱與索引
layerNames = net.getLayerNames()
lastLayerId = net.getLayerId(layerNames[-1])
lastLayer = net.getLayer(lastLayerId)
print(lastLayer.type)
# 檢測
net.setInput(cv.dnn.blobFromImage(image, size=(300, 300), swapRB=True, crop=False))
cvOut = net.forward()
for detection in cvOut[0,0,:,:]:
score = float(detection[2])
if score > 0.5:
left = detection[3]*w
top = detection[4]*h
right = detection[5]*w
bottom = detection[6]*h
# 繪制
cv.rectangle(image, (int(left), int(top)), (int(right), int(bottom)), (0, 255, 0), thickness=2)
cv.putText(image, "Pedestrian", (int(left), int(top-10)), cv.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 1)
cv.imshow('pedestrain_demo', image)
cv.imwrite("D:/Pedestrian.png", image)
cv.waitKey(0)
cv.destroyAllWindows()
交流群
歡迎加入公眾號讀者群一起和同行交流，目前有SLAM、三維視覺、傳感器、自動(dòng)駕駛、計(jì)算攝影、檢測、分割、識別、醫(yī)學(xué)影像、GAN、算法競賽等微信群（以后會(huì)逐漸細(xì)分），請掃描下面微信號加群，備注：”昵稱+學(xué)校/公司+研究方向“，例如：”張三 + 上海交大 + 視覺SLAM“。請按照格式備注，否則不予通過。添加成功后會(huì)根據(jù)研究方向邀請進(jìn)入相關(guān)微信群。請勿在群內(nèi)發(fā)送廣告，否則會(huì)請出群，謝謝理解~