1 背景
隨著Pytorch、TensorFlow等有效的框架被用來深度的學習開發(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