yolov5 onnxruntime c++ 推理 onnx

首先，您需要将Yolov5模型转换为ONNX格式。您可以使用PyTorch将模型转换为ONNX格式，然后使用ONNX Runtime C++ API加载和运行模型。

以下是一些步骤：

1. 安装PyTorch和ONNX Runtime

2. 使用PyTorch将Yolov5模型转换为ONNX格式。您可以使用以下代码：

import torch
import torchvision

# Load the model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# Export the model to ONNX format
input_shape = (1, 3, 640, 640)
torch.onnx.export(model, torch.randn(*input_shape), "yolov5s.onnx", opset_version=11)

3. 在C++中加载和运行模型。您可以使用以下代码：

#include <iostream>
#include <vector>
#include <chrono>
#include <opencv2/opencv.hpp>
#include "onnxruntime_cxx_api.h"

using namespace std;
using namespace cv;
using namespace std::chrono;
using namespace onnxruntime;

int main() {
    // Load the model
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");
    Ort::Session session(env, "yolov5s.onnx", session_options);

    // Get input and output names
    auto input_names = session.GetInputNames();
    auto output_names = session.GetOutputNames();

    // Create input tensor
    Ort::AllocatorWithDefaultOptions allocator;
    Ort::Value input_tensor(nullptr);
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    vector<int64_t> input_shape = {1, 3, 640, 640};
    input_tensor = Ort::Value::CreateTensor<float>(memory_info, reinterpret_cast<float*>(new float[input_shape[0] * input_shape[1] * input_shape[2] * input_shape[3]]), input_shape.data(), input_shape.size());

    // Load image
    Mat image = imread("test.jpg");
    cvtColor(image, image, COLOR_BGR2RGB);
    resize(image, image, Size(640, 640));
    float* input_data = input_tensor.GetTensorMutableData<float>();
    for (int i = 0; i < 640 * 640 * 3; i++) {
        input_data[i] = image.data[i] / 255.0;
    }

    // Run inference
    auto start = high_resolution_clock::now();
    vector<Ort::Value> output_tensors = session.Run(output_names, &input_names[0], &input_tensor, 1);
    auto end = high_resolution_clock::now();
    auto duration = duration_cast<milliseconds>(end - start);
    cout << "Inference time: " << duration.count() << " ms" << endl;

    // Get output tensor
    Ort::Value& output_tensor = output_tensors[0];
    float* output_data = output_tensor.GetTensorMutableData<float>();

    // Process output
    for (int i = 0; i < 25200; i++) {
        if (output_data[i * 6 + 4] > 0.5) {
            int x1 = output_data[i * 6 + 0] * 640;
            int y1 = output_data[i * 6 + 1] * 640;
            int x2 = output_data[i * 6 + 2] * 640;
            int y2 = output_data[i * 6 + 3] * 640;
            cout << "Object detected: " << output_data[i * 6 + 5] << " (" << x1 << ", " << y1 << ") (" << x2 << ", " << y2 << ")" << endl;
        }
    }

    return 0;
}

这个例子假设您有一张名为“test.jpg”的图像，它将被用作模型的输入。它还假设您的模型输出是一个大小为[1, 25200, 6]的张量，其中25200是预测的边界框数，6是每个边界框的属性数（左上角和右下角坐标，置信度和类别）。

请注意，这只是一个简单的例子，您需要根据您的模型和数据进行适当的修改。