0002-opencv dnn模块使用CUDA加速

opencv调用yolov4模型

可以参考0001这篇文章

安装opencv-contrib-python

如果要用GPU加速,一定要安装这个包。不然就会报错。

pip install opencv-contrib-python 

CUDA加速

添加以下两行代码就可以

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

案例脚本

import cv2 as cv
import numpy as np

classFile = "coco.names" 
with open(classFile, 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')
modelConf = 'coco.cfg' 
modelWeights = 'coco.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

参考链接

1. How to use OpenCV’s ‘dnn’ module with NVIDIA GPUs, CUDA, and cuDNN

你可能感兴趣的:(每日一个脚本强壮自己)