opencv-python dnn模块调用NVIDIA的GPU用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加速推理了,可以试下速度是否有提升!