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

這樣模型就加載完畢了,非常的簡單??梢钥闯鯫pencv的友好度還是非常好的,很容易轉換。

下面就是將圖片轉換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

?著作權歸作者所有,轉載或內容合作請聯系作者
平臺聲明:文章內容(如有圖片或視頻亦包括在內)由作者上傳并發布,文章內容僅代表作者本人觀點,簡書系信息發布平臺,僅提供信息存儲服務。

推薦閱讀更多精彩內容