如何用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

你可能感兴趣的:(如何用OpenCV加载Yolov5并使用CUDA加速)