如何用OpenCV加载Yolov5并使用CUDA加速

1 背景

随着PytorchTensorFlow等有效的框架被用来深度的学习开发,各种任务的模型也层出不穷。但是大多的部署往往依赖签名的两个框架,需要前面的两个框架大量的库。而且先前的Yolov3和Yolov4有官方直接支持,可以自接加载weights和cfg文件。部署起来相对来说就很简单,但是最新的Yolov5确实基于Pytorch版本的,这使用Opencv部署起来就稍微的麻烦了。可以这时候我们希望有没有一种方法能够使得模型的部署能够完全的摆脱框架,这样就能够做到模型的训练和模型的部署分开。而且模型的部署是一劳永逸的,部署的流程比较固定只要部署一次就可以,那么算法工程师就可以安心的在模型的算法研究上。

2 环境准备

话不多说,直接上工具。这里我们直接使用CUDA加速的Opencv来部署我们的算法模型(有加速版的为什么不用呢),这里首先需要编译出CUDA版本的Opencv具体可以参考我前面的一篇博文如何编译Opencv CUDA版本这里有详细的介绍编译的过程。在使用之前我们可以用如下的代码测试一下CUDA是否可以使用:

cv2.cuda.getCudaEnabledDeviceCount()

如果输出的值大于1,则证明我们的cuda可以使用。否则则证明CUDA版本的Opencv不能使用。

3 模型转换

这里主要使用将Yolov5模型转换ONNX模型,然后用Opencv来加载该模型。关于如何将Yolov5模型转换为ONNX请参考我的前一片博文,这里不再介绍。默认已经有转换好的模型了,下一步就直接去加载该模型了。

4 DNN模块加载模型

主要使用DNN模块去加载ONNX模型,然后去获得模型的推理结果。在调用模型之前我们需要使用Yolov5中已经实现的切片函数,这里直接使用就可以了:

def _make_grid(self, nx=20, ny=20):
        xv, yv = np.meshgrid(np.arange(ny), np.arange(nx))
        return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32)

首先我们调用DNN模块的readNetFromONN函数直接加载ONNX模型,该函数可以将ONNX模型转换为DNN模型了,通过下面的代码:

self.net=cv2.dnn.readNetFromONNX(".onnx")

这样模型就加载完毕了,非常的简单。可以看出Opencv的友好度还是非常好的,很容易转换。

下面就是将图片转换DNN模块能够读的格式,这里采用DNN模块中的blobFromImge模块:

srcImg=cv2.imread(img_path)
blob = cv2.dnn.blobFromImage(srcimg, 1 / 255.0, (self.inpWidth, self.inpHeight), [0, 0, 0], swapRB=True,crop=False)

将转换后的图片输入的DNN中也很简单,只需简单一行代码就可以:

self.net.setInput(blob)

最后我们要获得模型的输出即可,同样也是简单一行代码即可:

outs = self.net.forward(self.net.getUnconnectedOutLayersNames())[0]

模型输出这里已经获得了,不过该结果是一个整个的数组数据,我们还需要对结果处理一下才能进行下一步的处理。处理过程也可以参照Yolo的源码,这里就从中拿一段出来:

outs = 1 / (1 + np.exp(-outs))  ###定义sigmoid函数
row_ind = 0
for i in range(self.nl):
    h, w = int(self.inpHeight / self.stride[i]), int(self.inpWidth / self.stride[i])
    length = int(self.na * h * w)
    if self.grid[i].shape[2:4] != (h, w):
    self.grid[i] = self._make_grid(w, h)

    outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + np.tile(
    self.grid[i], (self.na, 1))) * int(self.stride[i])
    outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2) ** 2 * np.repeat(
    self.anchor_grid[i], h * w, axis=0)
    row_ind += length

5 输出结果的后处理

获得DNN模型的输出结果后,下一步就是对输出结果的后处理了。这一部分主要是对重新实现了Yolo的检测头的处理过程获得检测的物体类别以及检测框的位置,将检测结果还原到原图上去。

def postprocess(self, frame, outs):
        frameHeight = frame.shape[0]
        frameWidth = frame.shape[1]
        # 求缩放比例
        ratioh, ratiow = frameHeight / self.inpHeight, frameWidth / self.inpWidth
        # Scan through all the bounding boxes output from the network and keep only the
        # ones with high confidence scores. Assign the box's class label as the class with the highest score.
        classIds = []
        confidences = []
        boxes = []
        for detection in outs:
            scores = detection[5:]
            classId = np.argmax(scores)
            confidence = scores[classId]
            if confidence > self.confThreshold and detection[4] > self.objThreshold:
                center_x = int(detection[0] * ratiow)
                center_y = int(detection[1] * ratioh)
                width = int(detection[2] * ratiow)
                height = int(detection[3] * ratioh)
                left = int(center_x - width / 2)
                top = int(center_y - height / 2)
                classIds.append(classId)
                confidences.append(float(confidence) * float(detection[4]))
                boxes.append([left, top, width, height])

下面一步就是实现NMS算法了,这里可以自己去实现NMS算法也可以直接调用DNN模块的。实测两种方式实现的结果几乎没有差异,虽然Yolo源码中使用了DIOU的方式。下面直接调用就可以了:

# NMS非极大值抑制算法去掉重复的框
indices = cv2.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold)

这样我们就完成了整个过程,整体实现起来不是很复杂。具体代码仓库可以参考:
https://github.com/iwanggp/yolov5-opencv-pycpp-tensorrt

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 216,997评论 6 502
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 92,603评论 3 392
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 163,359评论 0 353
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 58,309评论 1 292
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 67,346评论 6 390
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 51,258评论 1 300
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 40,122评论 3 418
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,970评论 0 275
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 45,403评论 1 313
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,596评论 3 334
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,769评论 1 348
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 35,464评论 5 344
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 41,075评论 3 327
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,705评论 0 22
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,848评论 1 269
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 47,831评论 2 370
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 44,678评论 2 354

推荐阅读更多精彩内容