opencv-python dnn模块使用CUDA加速

dnn模块加载darknet模型

下面以使用opencv的dnn模块调用Darknet训练的yolo的模型。这里加载了我之前训练的一个人脸识别的模型。

weightsPath = "face.weights"
configPath = "face.cfg"
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)

安装opencv-contrib-python

如果系统环境下只安装了opencv-python，是不能够使用GPU加速的，无论是Intel的还是NVIDIA的都不行。我当时运行报下面这个错误：

	net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
AttributeError: module 'cv2.dnn' has no attribute 'DNN_TARGET_CUDA'

需要安装contrib这个扩展包

pip install opencv-contrib-python 

CUDA加速

添加下面两行代码：

net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

这里我一开始把Backend设置为DNN_BACKEND_CUDA，之后就报错：

cv2.error: OpenCV(4.2.0) opencv-python\opencv\modules\dnn\src\dnn.cpp:1304: error: (-215:Assertion failed) 
preferableBackend != DNN_BACKEND_OPENCV || 
preferableTarget == DNN_TARGET_CPU || preferableTarget == DNN_TARGET_OPENCL || 
preferableTarget == DNN_TARGET_OPENCL_FP16 in function 
'cv::dnn::dnn4_v20191202::Net::Impl::setUpNet'

所以这里要注意都设置成CUDA。

调用模型进行检测

import cv2 as cv
import numpy as np

classFile = "obj.names" 
with open(classFile, 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')
modelConf = 'obj.cfg' 
modelWeights = 'obj.weights'
net = cv.dnn.readNetFromDarknet(modelConf, modelWeights)
net.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)
cap = cv.VideoCapture(inputFile)
while(True):
      _,frame = cap.read()
     if np.shape(frame) != ():
             blob = cv.dnn.blobFromImage(frame, 1/255, (inpWidth, inpHeight), [0,0,0],1,crop=False)
             net.setInput(blob)
             outs = net.forward(getOutputsNames(net))
             frameExtract(frame, outs)
             cv.imshow("video", frame)
             k = cv.waitKey(1) & 0xFF
     else:
             print("Reinitialize capture device ", time.ctime())
             cap = cv.VideoCapture(inputFile)
             time.sleep(1)
             k = cv.waitKey(1) & 0xFF
     if k == 27:
            cv.destroyAllWindows()
            break

此代码参照下面这个链接：
链接: How to run OpenCV DNN on NVidia GPU.

设置好以后就可以调用cuda加速推理了，可以试下速度是否有提升！