【OpenCV】128 OpenCV DNN 直接调用tensorflow的导出模型

128 OpenCV DNN 直接调用tensorflow的导出模型

代码

import cv2 as cv

inference_pb = "faster_rcnn_resnet50_coco_2018_01_28/frozen_inference_graph.pb"
graph_text = "faster_rcnn_resnet50_coco_2018_01_28/graph.pbtxt"
obj_text = "faster_rcnn_resnet50_coco_2018_01_28/object_detection_classes_coco.txt"

# obj_name
objName = []
file = open(obj_text)
lines = file.readlines()
for line in lines:
    line = line.strip()
    objName.append(line)
print('objName: ', objName)


# load tensorflow model
net = cv.dnn.readNetFromTensorflow(inference_pb, graph_text)
image = cv.imread("../images/objects.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])
    objIndex = int(detection[1])
    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, "score:%.2f, %s" % (score, objName[objIndex]), (int(left), int(top)), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)

cv.imshow('faster-rcnn-demo', image)
cv.waitKey(0)
cv.destroyAllWindows()

实验结果

解释

OpenCV在DNN模块中支持直接调用tensorflow object detection训练导出的模型使用，支持的模型包括

• SSD
• Faster-RCNN
• Mask-RCNN
  三种经典的对象检测网络，这样就可以实现从tensorflow模型训练、导出模型、在OpenCV DNN调用模型网络实现自定义对象检测的技术链路，具有非常高的实用价值。以Faster-RCNN为例，模型下载地址如下

对于这些模型没有与之匹配的graph.pbtxt文件，OpenCV DNN模块提供python脚本来生成

---

所有内容均来源于贾志刚老师的知识星球——OpenCV研习社，本文为个人整理学习，已获得贾老师授权，有兴趣、有能力的可以加入贾老师的知识星球进行深入学习。