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

1 背景

隨著PytorchTensorFlow等有效的框架被用來深度的學習開發(fā)涧至,各種任務的模型也層出不窮治宣。但是大多的部署往往依賴簽名的兩個框架,需要前面的兩個框架大量的庫执赡。而且先前的Yolov3和Yolov4有官方直接支持,可以自接加載weights和cfg文件函筋。部署起來相對來說就很簡單沙合,但是最新的Yolov5確實基于Pytorch版本的,這使用Opencv部署起來就稍微的麻煩了跌帐∈仔福可以這時候我們希望有沒有一種方法能夠使得模型的部署能夠完全的擺脫框架,這樣就能夠做到模型的訓練和模型的部署分開谨敛。而且模型的部署是一勞永逸的究履,部署的流程比較固定只要部署一次就可以,那么算法工程師就可以安心的在模型的算法研究上脸狸。

2 環(huán)境準備

話不多說最仑,直接上工具。這里我們直接使用CUDA加速的Opencv來部署我們的算法模型(有加速版的為什么不用呢)炊甲,這里首先需要編譯出CUDA版本的Opencv具體可以參考我前面的一篇博文如何編譯Opencv CUDA版本這里有詳細的介紹編譯的過程泥彤。在使用之前我們可以用如下的代碼測試一下CUDA是否可以使用:

cv2.cuda.getCudaEnabledDeviceCount()

如果輸出的值大于1,則證明我們的cuda可以使用卿啡。否則則證明CUDA版本的Opencv不能使用吟吝。

3 模型轉(zhuǎn)換

這里主要使用將Yolov5模型轉(zhuǎn)換ONNX模型,然后用Opencv來加載該模型牵囤。關于如何將Yolov5模型轉(zhuǎn)換為ONNX請參考我的前一片博文爸黄,這里不再介紹。默認已經(jīng)有轉(zhuǎn)換好的模型了揭鳞,下一步就直接去加載該模型了炕贵。

4 DNN模塊加載模型

主要使用DNN模塊去加載ONNX模型,然后去獲得模型的推理結(jié)果野崇。在調(diào)用模型之前我們需要使用Yolov5中已經(jīng)實現(xiàn)的切片函數(shù)称开,這里直接使用就可以了:

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)

首先我們調(diào)用DNN模塊的readNetFromONN函數(shù)直接加載ONNX模型,該函數(shù)可以將ONNX模型轉(zhuǎn)換為DNN模型了乓梨,通過下面的代碼:

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

這樣模型就加載完畢了鳖轰,非常的簡單》龆疲可以看出Opencv的友好度還是非常好的蕴侣,很容易轉(zhuǎn)換。

下面就是將圖片轉(zhuǎn)換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)

將轉(zhuǎn)換后的圖片輸入的DNN中也很簡單昆雀,只需簡單一行代碼就可以:

self.net.setInput(blob)

最后我們要獲得模型的輸出即可,同樣也是簡單一行代碼即可:

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

模型輸出這里已經(jīng)獲得了蝠筑,不過該結(jié)果是一個整個的數(shù)組數(shù)據(jù)狞膘,我們還需要對結(jié)果處理一下才能進行下一步的處理。處理過程也可以參照Yolo的源碼什乙,這里就從中拿一段出來:

outs = 1 / (1 + np.exp(-outs))  ###定義sigmoid函數(shù)
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 輸出結(jié)果的后處理

獲得DNN模型的輸出結(jié)果后挽封,下一步就是對輸出結(jié)果的后處理了。這一部分主要是對重新實現(xiàn)了Yolo的檢測頭的處理過程獲得檢測的物體類別以及檢測框的位置臣镣,將檢測結(jié)果還原到原圖上去辅愿。

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])

下面一步就是實現(xiàn)NMS算法了,這里可以自己去實現(xiàn)NMS算法也可以直接調(diào)用DNN模塊的退疫。實測兩種方式實現(xiàn)的結(jié)果幾乎沒有差異渠缕,雖然Yolo源碼中使用了DIOU的方式。下面直接調(diào)用就可以了:

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

這樣我們就完成了整個過程褒繁,整體實現(xiàn)起來不是很復雜亦鳞。具體代碼倉庫可以參考:
https://github.com/iwanggp/yolov5-opencv-pycpp-tensorrt

?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
  • 序言:七十年代末,一起剝皮案震驚了整個濱河市棒坏,隨后出現(xiàn)的幾起案子燕差,更是在濱河造成了極大的恐慌,老刑警劉巖坝冕,帶你破解...
    沈念sama閱讀 223,207評論 6 521
  • 序言:濱河連續(xù)發(fā)生了三起死亡事件徒探,死亡現(xiàn)場離奇詭異,居然都是意外死亡喂窟,警方通過查閱死者的電腦和手機测暗,發(fā)現(xiàn)死者居然都...
    沈念sama閱讀 95,455評論 3 400
  • 文/潘曉璐 我一進店門央串,熙熙樓的掌柜王于貴愁眉苦臉地迎上來,“玉大人碗啄,你說我怎么就攤上這事质和。” “怎么了稚字?”我有些...
    開封第一講書人閱讀 170,031評論 0 366
  • 文/不壞的土叔 我叫張陵饲宿,是天一觀的道長。 經(jīng)常有香客問我胆描,道長瘫想,這世上最難降的妖魔是什么? 我笑而不...
    開封第一講書人閱讀 60,334評論 1 300
  • 正文 為了忘掉前任昌讲,我火速辦了婚禮国夜,結(jié)果婚禮上,老公的妹妹穿的比我還像新娘短绸。我一直安慰自己支竹,他們只是感情好,可當我...
    茶點故事閱讀 69,322評論 6 398
  • 文/花漫 我一把揭開白布鸠按。 她就那樣靜靜地躺著礼搁,像睡著了一般。 火紅的嫁衣襯著肌膚如雪目尖。 梳的紋絲不亂的頭發(fā)上馒吴,一...
    開封第一講書人閱讀 52,895評論 1 314
  • 那天,我揣著相機與錄音瑟曲,去河邊找鬼饮戳。 笑死,一個胖子當著我的面吹牛洞拨,可吹牛的內(nèi)容都是我干的扯罐。 我是一名探鬼主播,決...
    沈念sama閱讀 41,300評論 3 424
  • 文/蒼蘭香墨 我猛地睜開眼烦衣,長吁一口氣:“原來是場噩夢啊……” “哼歹河!你這毒婦竟也來了?” 一聲冷哼從身側(cè)響起花吟,我...
    開封第一講書人閱讀 40,264評論 0 277
  • 序言:老撾萬榮一對情侶失蹤秸歧,失蹤者是張志新(化名)和其女友劉穎,沒想到半個月后衅澈,有當?shù)厝嗽跇淞掷锇l(fā)現(xiàn)了一具尸體键菱,經(jīng)...
    沈念sama閱讀 46,784評論 1 321
  • 正文 獨居荒郊野嶺守林人離奇死亡,尸身上長有42處帶血的膿包…… 初始之章·張勛 以下內(nèi)容為張勛視角 年9月15日...
    茶點故事閱讀 38,870評論 3 343
  • 正文 我和宋清朗相戀三年今布,在試婚紗的時候發(fā)現(xiàn)自己被綠了经备。 大學時的朋友給我發(fā)了我未婚夫和他白月光在一起吃飯的照片拭抬。...
    茶點故事閱讀 40,989評論 1 354
  • 序言:一個原本活蹦亂跳的男人離奇死亡,死狀恐怖侵蒙,靈堂內(nèi)的尸體忽然破棺而出玖喘,到底是詐尸還是另有隱情,我是刑警寧澤蘑志,帶...
    沈念sama閱讀 36,649評論 5 351
  • 正文 年R本政府宣布,位于F島的核電站贬派,受9級特大地震影響急但,放射性物質(zhì)發(fā)生泄漏。R本人自食惡果不足惜搞乏,卻給世界環(huán)境...
    茶點故事閱讀 42,331評論 3 336
  • 文/蒙蒙 一波桩、第九天 我趴在偏房一處隱蔽的房頂上張望。 院中可真熱鬧请敦,春花似錦镐躲、人聲如沸。這莊子的主人今日做“春日...
    開封第一講書人閱讀 32,814評論 0 25
  • 文/蒼蘭香墨 我抬頭看了看天上的太陽。三九已至匣椰,卻和暖如春裆熙,著一層夾襖步出監(jiān)牢的瞬間,已是汗流浹背禽笑。 一陣腳步聲響...
    開封第一講書人閱讀 33,940評論 1 275
  • 我被黑心中介騙來泰國打工入录, 沒想到剛下飛機就差點兒被人妖公主榨干…… 1. 我叫王不留,地道東北人佳镜。 一個月前我還...
    沈念sama閱讀 49,452評論 3 379
  • 正文 我出身青樓僚稿,卻偏偏與公主長得像,于是被迫代替她去往敵國和親蟀伸。 傳聞我的和親對象是個殘疾皇子蚀同,可洞房花燭夜當晚...
    茶點故事閱讀 45,995評論 2 361

推薦閱讀更多精彩內(nèi)容