在C++中使用OpenVINO工具包部署YOLOv5-Seg模型

來源：小白學視覺

YOLOv5兼具速度和精度，工程化做的特別好，Git clone到本地即可在自己的數(shù)據(jù)集上實現(xiàn)目標檢測任務的訓練和推理，在產(chǎn)業(yè)界中應用廣泛。開源社區(qū)對YOLOv5支持實例分割的呼聲高漲，YOLOv5在v7.0中正式官宣支持實例分割。

本文主要介紹在C++中使用OpenVINO工具包部署YOLOv5-Seg模型，主要步驟有：

配置OpenVINO C++開發(fā)環(huán)境

下載并轉換YOLOv5-Seg預訓練模型

使用OpenVINO Runtime C++ API編寫推理程序

下面，本文將依次詳述

1.1 配置OpenVINO C++開發(fā)環(huán)境

配置OpenVINO C++開發(fā)環(huán)境的詳細步驟，請參考《在Windows中基于Visual Studio配置OpenVINO C++開發(fā)環(huán)境》。

1.2 下載并轉換YOLOv5預訓練模型

下載并轉換YOLOv5-seg預訓練模型的詳細步驟，請參考：https://mp.weixin.qq.com/s/K3wP5YLAU4p5jsdiMYjuMg，本文所使用的OpenVINO是2022.3 LTS版。

首先，運行命令獲得 yolov5s-seg ONNX 格式模型：yolov5s-seg.onnx：

python export.py --weights yolov5s-seg.pt --include onnx

然后運行命令獲得yolov5s-seg IR格式模型：yolov5s-seg.xml和yolov5s-seg.bin，如下圖所示

mo -m yolov5s-seg.onnx --compress_to_fp16

1.3 使用OpenVINO Runtime C++ API編寫推理程序

一個端到端的AI推理程序，主要包含五個典型的處理流程：

采集圖像&圖像解碼

圖像數(shù)據(jù)預處理

AI推理計算

對推理結果進行后處理

將處理后的結果集成到業(yè)務流程

圖 1-2 端到端的AI推理程序處理流程

1.3.1 采集圖像&圖像解碼

OpenCV提供imread()函數(shù)將圖像文件載入內(nèi)存，

Mat cv::imread (const String &filename, int flags=IMREAD_COLOR)

若是從視頻流(例如，視頻文件、網(wǎng)絡攝像頭、3D攝像頭(Realsense)等)中，一幀一幀讀取圖像數(shù)據(jù)到內(nèi)存，則使用cv::VideoCapture類，對應范例代碼請參考OpenCV官方范例代碼：https://github.com/opencv/opencv/tree/4.x/samples/cpp。

圖 1-3 從視頻流讀取圖像幀范例

1.3.2 YOLOv5-Seg模型的圖像預處理

YOLOv5-Seg模型構架是在YOLOv5模型構架基礎上，增加了一個叫“Proto”的小型卷積神經(jīng)網(wǎng)絡，用于輸出檢測對象掩碼(Mask)，如下圖所示：

圖 1-4 YOLOv5-Seg模型輸出的代碼定義

詳細參看：https://github.com/ultralytics/yolov5/blob/master/models/yolo.py#L92

由此可知，YOLOv5-Seg模型對數(shù)據(jù)預處理的要求跟YOLOv5模型一模一樣，YOLOv5-Seg模型的預處理代碼可以復用YOLOv5模型的C++預處理代碼。

另外，從代碼可以看出YOLOv5-Seg模型的輸出有兩個張量，一個張量輸出檢測結果，一個張量輸出proto，其形狀可以用Netron打開yolov5-seg.onnx查知，如下圖所示。

圖 1-5 YOLOv5-Seg模型的輸入和輸出

“output0”是檢測輸出，第一個維度表示batch size，第二個維度表示25200條輸出，第三個維度表示有117個字段，其中前85個字段(0~84)表示：cx、cy、w、h、confidence和80個類別分數(shù)，后32個字段與”output1”做矩陣乘法，可以獲得尺寸為160x160的檢測目標的掩碼(mask)。

1.3.3 執(zhí)行AI推理計算

基于OpenVINO Runtime C++ API實現(xiàn)AI推理計算主要有兩種方式：一種是同步推理方式，一種是異步推理方式，本文主要介紹同步推理方式。

主要步驟有：

初始化Core類：ov::Core core;

編譯模型：core.compile_model()

創(chuàng)建推理請求infer_request：compiled_model.create_infer_request()

讀取圖像數(shù)據(jù)并做預處理：letterbox()

將預處理后的blob數(shù)據(jù)傳入模型輸入節(jié)點：infer_request.set_input_tensor()

調用infer()方法執(zhí)行推理計算：infer_request.infer()

獲得推理結果：infer_request.get_output_tensor()

基于OpenVINO Runtime C++API的同步推理代碼如下所示：

 // -------- Step 1. Initialize OpenVINO Runtime Core --------
    ov::Core core;
    // -------- Step 2. Compile the Model --------
    auto compiled_model = core.compile_model(model_file, "GPU.1"); //GPU.1 is dGPU A770
    // -------- Step 3. Create an Inference Request --------
    ov::InferRequest infer_request = compiled_model.create_infer_request();
    // -------- Step 4. Read a picture file and do the preprocess --------
    cv::Mat img = cv::imread(image_file); //Load a picture into memory
    std::vector paddings(3);       //scale, half_h, half_w
    cv::Mat resized_img = letterbox(img, paddings); //resize to (640,640) by letterbox
    // BGR->RGB, u8(0-255)->f32(0.0-1.0), HWC->NCHW
    cv::Mat blob = cv::blobFromImage(resized_img, 1 / 255.0, cv::Size(640, 640), cv::Scalar(0, 0, 0), true);
    // -------- Step 5. Feed the blob into the input node of YOLOv5 -------
    // Get input port for model with one input
    auto input_port = compiled_model.input();
    // Create tensor from external memory
    ov::Tensor input_tensor(input_port.get_element_type(), input_port.get_shape(), blob.ptr(0));
    // Set input tensor for model with one input
    infer_request.set_input_tensor(input_tensor);
    // -------- Step 6. Start inference --------
    infer_request.infer();
    // -------- Step 7. Get the inference result --------
    auto detect = infer_request.get_output_tensor(0);
    auto detect_shape = detect.get_shape();
    std::cout << "The shape of Detection tensor:"<< detect_shape << std::endl;
    auto proto = infer_request.get_output_tensor(1);
    auto proto_shape = proto.get_shape();
std::cout << "The shape of Proto tensor:" << proto_shape << std::endl;

1.3.4 推理結果進行后處理

后處理工作主要是從”detect ”輸出張量中拆解出檢測框的位置和類別信息，并用cv::NMSBoxes()過濾掉多于的檢測框；從”detect ”輸出張量的后32個字段與”proto”輸出張量做矩陣乘法，獲得每個檢測目標的形狀為160x160的掩碼輸出，最后將160x160的掩碼映射回原始圖像完成所有后處理工作。

1.4 總結

配置OpenVINO C++開發(fā)環(huán)境后，可以直接編譯運行yolov5seg_openvino_dGPU.cpp，結果如下圖所示。使用OpenVINO Runtime C++ API函數(shù)開發(fā)YOLOv5推理程序，簡單方便，并可以任意部署在英特爾CPU、集成顯卡和獨立顯卡上。