使用openvino进行onnx的模型推理过程

前言:

openvino是Intel开发的基于intel设备的推理引擎,可以利用CPU发挥最好的性能,还能使用到新款CPU所提供的NN算力。

cpp

int main(){
    inference();
    return 0;
}

main函数还是只执行了一个inference

使用openvino进行onnx的模型推理过程_第1张图片

这里的link可以看到是使用了intel的tbb库,这个我们在之前讲到过TBB的并行模式。

Inference:

void inference(){

    size_t input_batch = 1;
    size_t input_channel = 3;
    size_t input_height = 640;
    size_t input_width = 640;

    ov::Core core;
    auto model = core.compile_model("yolov5s.onnx");
    auto iq = model.create_infer_request();

这里的ov就是openvino的意思,这里是先定义了一个core,然后将onnx文件compile_model进去得到了一个model,然后创立了infer_request,是一个推理的请求

    auto input = iq.get_input_tensor(0);
    auto output = iq.get_output_tensor(0);
    input.set_shape({input_batch, input_channel, input_height, input_width});

    float* input_data_host = input.data();

通过Iq拿到了intput和output的tensor

    auto input = iq.get_input_tensor(0);
    auto output = iq.get_output_tensor(0);
    input.set_shape({input_batch, input_channel, input_height, input_width});

    float* input_data_host = input.data();

使用openvino进行onnx的模型推理过程_第2张图片

因为yolov5是一个动态的shape,所以需要set_shape来确定shape,也就是我们之前所设定的1*3*640*640。

之后将其分配内存。

预处理:

    ///
    // letter box
    auto image = cv::imread("car.jpg");
    // 通过双线性插值对图像进行resize
    float scale_x = input_width / (float)image.cols;
    float scale_y = input_height / (float)image.rows;
    float scale = std::min(scale_x, scale_y);
    float i2d[6], d2i[6];
    // resize图像,源图像和目标图像几何中心的对齐
    i2d[0] = scale;  i2d[1] = 0;  i2d[2] = (-scale * image.cols + input_width + scale  - 1) * 0.5;
    i2d[3] = 0;  i2d[4] = scale;  i2d[5] = (-scale * image.rows + input_height + scale - 1) * 0.5;

    cv::Mat m2x3_i2d(2, 3, CV_32F, i2d);  // image to dst(network), 2x3 matrix
    cv::Mat m2x3_d2i(2, 3, CV_32F, d2i);  // dst to image, 2x3 matrix
    cv::invertAffineTransform(m2x3_i2d, m2x3_d2i);  // 计算一个反仿射变换

    cv::Mat input_image(input_height, input_width, CV_8UC3);
    cv::warpAffine(image, input_image, m2x3_i2d, input_image.size(), cv::INTER_LINEAR, cv::BORDER_CONSTANT, cv::Scalar::all(114));  // 对图像做平移缩放旋转变换,可逆
    cv::imwrite("input-image.jpg", input_image);

    int image_area = input_image.cols * input_image.rows;
    unsigned char* pimage = input_image.data;
    float* phost_b = input_data_host + image_area * 0;
    float* phost_g = input_data_host + image_area * 1;
    float* phost_r = input_data_host + image_area * 2;
    for(int i = 0; i < image_area; ++i, pimage += 3){
        // 注意这里的顺序rgb调换了
        *phost_r++ = pimage[0] / 255.0f;
        *phost_g++ = pimage[1] / 255.0f;
        *phost_b++ = pimage[2] / 255.0f;
    }
    ///

推理:

将image塞到input_image_host里后作推理,只有一句话:

   iq.infer();

推理过后output这个tensor会直接获取到data

    int output_numbox = output.get_shape()[1];
    int output_numprob = output.get_shape()[2];
    int num_classes = output_numprob - 5;
    float* output_data_host = output.data();

后处理:

// decode box:从不同尺度下的预测狂还原到原输入图上(包括:预测框,类被概率,置信度)
    vector> bboxes;
    float confidence_threshold = 0.25;
    float nms_threshold = 0.5;
    for(int i = 0; i < output_numbox; ++i){
        float* ptr = output_data_host + i * output_numprob;
        float objness = ptr[4];
        if(objness < confidence_threshold)
            continue;

        float* pclass = ptr + 5;
        int label     = std::max_element(pclass, pclass + num_classes) - pclass;
        float prob    = pclass[label];
        float confidence = prob * objness;
        if(confidence < confidence_threshold)
            continue;

        // 中心点、宽、高
        float cx     = ptr[0];
        float cy     = ptr[1];
        float width  = ptr[2];
        float height = ptr[3];

        // 预测框
        float left   = cx - width * 0.5;
        float top    = cy - height * 0.5;
        float right  = cx + width * 0.5;
        float bottom = cy + height * 0.5;

        // 对应图上的位置
        float image_base_left   = d2i[0] * left   + d2i[2];
        float image_base_right  = d2i[0] * right  + d2i[2];
        float image_base_top    = d2i[0] * top    + d2i[5];
        float image_base_bottom = d2i[0] * bottom + d2i[5];
        bboxes.push_back({image_base_left, image_base_top, image_base_right, image_base_bottom, (float)label, confidence});
    }
    printf("decoded bboxes.size = %d\n", bboxes.size());

    // nms非极大抑制
    std::sort(bboxes.begin(), bboxes.end(), [](vector& a, vector& b){return a[5] > b[5];});
    std::vector remove_flags(bboxes.size());
    std::vector> box_result;
    box_result.reserve(bboxes.size());

    auto iou = [](const vector& a, const vector& b){
        float cross_left   = std::max(a[0], b[0]);
        float cross_top    = std::max(a[1], b[1]);
        float cross_right  = std::min(a[2], b[2]);
        float cross_bottom = std::min(a[3], b[3]);

        float cross_area = std::max(0.0f, cross_right - cross_left) * std::max(0.0f, cross_bottom - cross_top);
        float union_area = std::max(0.0f, a[2] - a[0]) * std::max(0.0f, a[3] - a[1]) 
                         + std::max(0.0f, b[2] - b[0]) * std::max(0.0f, b[3] - b[1]) - cross_area;
        if(cross_area == 0 || union_area == 0) return 0.0f;
        return cross_area / union_area;
    };

    for(int i = 0; i < bboxes.size(); ++i){
        if(remove_flags[i]) continue;

        auto& ibox = bboxes[i];
        box_result.emplace_back(ibox);
        for(int j = i + 1; j < bboxes.size(); ++j){
            if(remove_flags[j]) continue;

            auto& jbox = bboxes[j];
            if(ibox[4] == jbox[4]){
                // class matched
                if(iou(ibox, jbox) >= nms_threshold)
                    remove_flags[j] = true;
            }
        }
    }
    printf("box_result.size = %d\n", box_result.size());

    for(int i = 0; i < box_result.size(); ++i){
        auto& ibox = box_result[i];
        float left = ibox[0];
        float top = ibox[1];
        float right = ibox[2];
        float bottom = ibox[3];
        int class_label = ibox[4];
        float confidence = ibox[5];
        cv::Scalar color;
        tie(color[0], color[1], color[2]) = random_color(class_label);
        cv::rectangle(image, cv::Point(left, top), cv::Point(right, bottom), color, 3);

        auto name      = cocolabels[class_label];
        auto caption   = cv::format("%s %.2f", name, confidence);
        int text_width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10;
        cv::rectangle(image, cv::Point(left-3, top-33), cv::Point(left + text_width, top), color, -1);
        cv::putText(image, caption, cv::Point(left, top-5), 0, 1, cv::Scalar::all(0), 2, 16);
    }
    cv::imwrite("image-draw.jpg", image);
}

总结:

  1. - 通过model = core.compile_model(xx.onnx)编译模型
  2. - 通过iq = model.create_infer_request()创建推理请求
  3. - input = iq.get_input_tensor(0);获取输入的tensor
  4. - output = iq.get_output_tensor(0);获取输出的tensor
  5. - input.set_shape({input_batch, input_channel, input_height, input_width});配置输入大小,因为是动态batch,需要先设置大小,此时会分配空间
  6. - input_data_host = input.data();获取输入指针,必须set shape后才能获取数据指针,否则会存储空间没分配而异常
  7. - 把图像预处理并储存到 input_data_host
  8. - iq.infer() 执行推理步骤
  9. - output_data_host = output.data();通过output拿到推理后的输出
  10. - 对output data进行解码得到最后的输出框。

你可能感兴趣的:(openvino,trtONNX,c++)