YOLOv5 在最新 OpenVINO 2021R02 版本的部署與程式碼演示詳解

YOLOv5 在 OpenVINO 上的部署，網上有很多 python 版本的程式碼，但是基本都有個很內傷的問題，就是還在用 pytorch 的一些庫做解析，C++的程式碼有個更大的內傷就是自定義解析解釋的不是很清楚，所以本人閱讀 YOLOv5 的 pytorch 程式碼推理部分，從原始的三個輸出層解析實現了 boxes, classes, nms 等關鍵 C++ 程式碼輸出，實現了純 OpenVINO+OpenCV 版本的 YOLOv5s 模型推理的程式碼演示。下面就是詳細的系統環境與各個部分解釋，以及程式碼實現與演示影像。

 

系統版本資訊與依賴

 

 

Window 10 64bit

 

 

Pytorch1.7+CUDA10.0

 

 

Python3.8.5

 

 

VS2015

 

 

OpenVINO_2021.2.185

 

 

01

 

YOLOv5 下載與測試執行

YOLOv5 是第二個非官方的 YOLO 物件檢測版本，也是第一個 Pytorch 實現的 YOLO 物件檢測版本。Github 地址如下：

需要克隆到本地

git clone .git

然後執行

pip install -r requirements.txt

安裝所有依賴。

 

最後執行一段影片或者影像完成測試

python detect.py --source D:\images\video\SungEun.avi --weights yolov5s.pt --conf 0.25

02

 

模型轉換

模型轉換主要是把原始的 YOLOv5 的 pytorch 模型檔案轉換為通用的開放模型格式 ONNX 與 OpenVIN 特有的檔案格式IR(*.xml與*.bin)。

 

OpenVINO 從 2020R02 以後版本開始支援直接讀取 ONNX 格式檔案，所以我們既可以透過指令碼直接匯出 onnx 格式檔案，直接給 OpenVINO 呼叫，也 可以對得到 ONNX 檔案透過 OpenVINO 的模型轉換指令碼做進一步轉換生成IR中間格式(*.bin檔案與*.xml檔案)。

 

匯出ONNX格式檔案的指令碼

Pytorch 的 YOLOv5 專案本身已經提供了轉換指令碼，命令列執行方式如下：

# export at 640x640 with batch size 1python models/export.py --weights yolov5s.pt --img 640 --batch 1

然後迴圈每個輸出層，解析每個特徵點對應的3個框與相關資料。由於在匯出的時候 ONNX 格式檔案時模型的推理得到的三個輸出層原始結果，所以還需要對每個資料先完成 sigmoid 歸一化，然後再計算相關值，這部分的程式碼實現我參考了專案中的 yolo.py 中的 Detection 部分，得到初始每個物件的檢測框之後，採用 OpenVINO 中自帶非最大抑制函式，完成非最大抑制，就得到了最終的預測框，然後繪製顯示。