feat: 实现 ONNX Runtime CGO 绑定和 OCR/Math 推理逻辑

2026-07-16 21:14:34 +00:00
parent 5dacb61122
commit 7afefe5e9f
6 changed files with 1094 additions and 85 deletions
@@ -0,0 +1,424 @@
+#include <onnxruntime_c_api.h>
+#include <vector>
+#include <string>
+#include <cstring>
+#include <memory>
+
+static const OrtApi* g_ort = nullptr;
+static std::string last_error;
+
+// 初始化 ONNX Runtime
+extern "C" int onnx_init() {
+    g_ort = OrtGetApiBase()->GetApi(ORT_API_VERSION);
+    if (!g_ort) {
+        last_error = "Failed to get ONNX Runtime API";
+        return -1;
+    }
+    return 0;
+}
+
+// ONNX Session 结构
+typedef struct {
+    OrtSession* session;
+    OrtSessionOptions* session_options;
+    OrtMemoryInfo* memory_info;
+    std::vector<std::string> input_names;
+    std::vector<std::string> output_names;
+    std::vector<const char*> input_name_ptrs;
+    std::vector<const char*> output_name_ptrs;
+} OnnxSession;
+
+// 创建 Session
+extern "C" OnnxSession* onnx_create_session(const char* model_path) {
+    if (!g_ort) {
+        onnx_init();
+    }
+
+    auto* sess = new OnnxSession();
+    
+    // 创建 session options
+    OrtStatus* status = g_ort->CreateSessionOptions(&sess->session_options);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        delete sess;
+        return nullptr;
+    }
+
+    // 设置 CPU 线程数
+    g_ort->SetIntraOpNumThreads(sess->session_options, 4);
+    g_ort->SetSessionGraphOptimizationLevel(sess->session_options, GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
+
+    // 创建 memory info
+    status = g_ort->CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &sess->memory_info);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        g_ort->ReleaseSessionOptions(sess->session_options);
+        delete sess;
+        return nullptr;
+    }
+
+    // 创建 session
+    status = g_ort->CreateSession(model_path, sess->session_options, &sess->session);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        g_ort->ReleaseMemoryInfo(sess->memory_info);
+        g_ort->ReleaseSessionOptions(sess->session_options);
+        delete sess;
+        return nullptr;
+    }
+
+    // 获取输入输出名称
+    OrtAllocator* allocator = nullptr;
+    status = g_ort->GetAllocatorWithDefaultOptions(&allocator);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+    } else {
+        // 获取输入数量和名称
+        size_t num_inputs = 0;
+        g_ort->SessionGetInputCount(sess->session, &num_inputs);
+        sess->input_names.resize(num_inputs);
+        sess->input_name_ptrs.resize(num_inputs);
+        
+        for (size_t i = 0; i < num_inputs; i++) {
+            char* name = nullptr;
+            g_ort->SessionGetInputName(sess->session, i, allocator, &name);
+            sess->input_names[i] = name;
+            sess->input_name_ptrs[i] = sess->input_names[i].c_str();
+            g_ort->AllocatorFree(allocator, name);
+        }
+
+        // 获取输出数量和名称
+        size_t num_outputs = 0;
+        g_ort->SessionGetOutputCount(sess->session, &num_outputs);
+        sess->output_names.resize(num_outputs);
+        sess->output_name_ptrs.resize(num_outputs);
+        
+        for (size_t i = 0; i < num_outputs; i++) {
+            char* name = nullptr;
+            g_ort->SessionGetOutputName(sess->session, i, allocator, &name);
+            sess->output_names[i] = name;
+            sess->output_name_ptrs[i] = sess->output_names[i].c_str();
+            g_ort->AllocatorFree(allocator, name);
+        }
+    }
+
+    return sess;
+}
+
+// 销毁 Session
+extern "C" void onnx_destroy_session(OnnxSession* sess) {
+    if (sess) {
+        if (sess->session) g_ort->ReleaseSession(sess->session);
+        if (sess->memory_info) g_ort->ReleaseMemoryInfo(sess->memory_info);
+        if (sess->session_options) g_ort->ReleaseSessionOptions(sess->session_options);
+        delete sess;
+    }
+}
+
+// 运行推理 - 支持动态输入
+extern "C" float* onnx_run_float(
+    OnnxSession* sess,
+    const float* input_data,
+    int input_size,
+    const int64_t* input_dims,
+    int input_dim_count,
+    int* output_size
+) {
+    if (!sess || !sess->session || !input_data) {
+        last_error = "Invalid session or input";
+        return nullptr;
+    }
+
+    OrtStatus* status = nullptr;
+
+    // 创建输入 tensor
+    OrtValue* input_tensor = nullptr;
+    status = g_ort->CreateTensorWithDataAsOrtValue(
+        sess->memory_info,
+        (void*)input_data,
+        input_size * sizeof(float),
+        input_dims,
+        input_dim_count,
+        ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
+        &input_tensor
+    );
+
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        return nullptr;
+    }
+
+    // 运行推理
+    OrtValue* output_tensor = nullptr;
+    status = g_ort->Run(
+        sess->session,
+        nullptr,
+        sess->input_name_ptrs.data(),
+        &input_tensor,
+        1,
+        sess->output_name_ptrs.data(),
+        1,
+        &output_tensor
+    );
+
+    g_ort->ReleaseValue(input_tensor);
+
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        return nullptr;
+    }
+
+    // 获取输出数据
+    float* output_data = nullptr;
+    status = g_ort->GetTensorMutableData(output_tensor, (void**)&output_data);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        g_ort->ReleaseValue(output_tensor);
+        return nullptr;
+    }
+
+    // 获取输出大小
+    OrtTensorTypeAndShapeInfo* type_info = nullptr;
+    status = g_ort->GetTensorTypeAndShape(output_tensor, &type_info);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        g_ort->ReleaseValue(output_tensor);
+        return nullptr;
+    }
+
+    size_t element_count = 0;
+    g_ort->GetTensorShapeElementCount(type_info, &element_count);
+    *output_size = (int)element_count;
+
+    // 复制输出数据
+    float* result = (float*)malloc(element_count * sizeof(float));
+    memcpy(result, output_data, element_count * sizeof(float));
+
+    g_ort->ReleaseTensorTypeAndShapeInfo(type_info);
+    g_ort->ReleaseValue(output_tensor);
+
+    return result;
+}
+
+// 获取输入形状
+extern "C" int onnx_get_input_shape(
+    OnnxSession* sess,
+    int input_index,
+    int64_t* dims,
+    int max_dims
+) {
+    if (!sess || !sess->session) {
+        last_error = "Invalid session";
+        return -1;
+    }
+
+    OrtTypeInfo* type_info = nullptr;
+    OrtStatus* status = g_ort->SessionGetInputTypeInfo(sess->session, input_index, &type_info);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        return -1;
+    }
+
+    OrtTensorTypeAndShapeInfo* tensor_info = nullptr;
+    status = g_ort->CastTypeInfoToTensorInfo(type_info, &tensor_info);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        g_ort->ReleaseTypeInfo(type_info);
+        return -1;
+    }
+
+    size_t dim_count = 0;
+    g_ort->GetDimensionsCount(tensor_info, &dim_count);
+    
+    if ((int)dim_count > max_dims) {
+        dim_count = max_dims;
+    }
+
+    g_ort->GetDimensions(tensor_info, dims, dim_count);
+    
+    g_ort->ReleaseTensorTypeAndShapeInfo(tensor_info);
+    g_ort->ReleaseTypeInfo(type_info);
+
+    return (int)dim_count;
+}
+
+// 获取输出形状
+extern "C" int onnx_get_output_shape(
+    OnnxSession* sess,
+    int output_index,
+    int64_t* dims,
+    int max_dims
+) {
+    if (!sess || !sess->session) {
+        last_error = "Invalid session";
+        return -1;
+    }
+
+    OrtTypeInfo* type_info = nullptr;
+    OrtStatus* status = g_ort->SessionGetOutputTypeInfo(sess->session, output_index, &type_info);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        return -1;
+    }
+
+    OrtTensorTypeAndShapeInfo* tensor_info = nullptr;
+    status = g_ort->CastTypeInfoToTensorInfo(type_info, &tensor_info);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        g_ort->ReleaseTypeInfo(type_info);
+        return -1;
+    }
+
+    size_t dim_count = 0;
+    g_ort->GetDimensionsCount(tensor_info, &dim_count);
+    
+    if ((int)dim_count > max_dims) {
+        dim_count = max_dims;
+    }
+
+    g_ort->GetDimensions(tensor_info, dims, dim_count);
+    
+    g_ort->ReleaseTensorTypeAndShapeInfo(tensor_info);
+    g_ort->ReleaseTypeInfo(type_info);
+
+    return (int)dim_count;
+}
+
+// 释放内存
+extern "C" void onnx_free(void* ptr) {
+    if (ptr) {
+        free(ptr);
+    }
+}
+
+// 获取错误信息
+extern "C" const char* onnx_get_last_error() {
+    return last_error.c_str();
+}
+
+// 双输入推理 (用于 Siamese 网络)
+extern "C" float* onnx_run_dual_input(
+    OnnxSession* sess,
+    const float* input1_data,
+    int input1_size,
+    const int64_t* input1_dims,
+    int input1_dim_count,
+    const float* input2_data,
+    int input2_size,
+    const int64_t* input2_dims,
+    int input2_dim_count,
+    int* output_size
+) {
+    if (!sess || !sess->session || !input1_data || !input2_data) {
+        last_error = "Invalid session or inputs";
+        return nullptr;
+    }
+
+    OrtStatus* status = nullptr;
+
+    // 创建输入 tensor 1
+    OrtValue* input_tensor1 = nullptr;
+    status = g_ort->CreateTensorWithDataAsOrtValue(
+        sess->memory_info,
+        (void*)input1_data,
+        input1_size * sizeof(float),
+        input1_dims,
+        input1_dim_count,
+        ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
+        &input_tensor1
+    );
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        return nullptr;
+    }
+
+    // 创建输入 tensor 2
+    OrtValue* input_tensor2 = nullptr;
+    status = g_ort->CreateTensorWithDataAsOrtValue(
+        sess->memory_info,
+        (void*)input2_data,
+        input2_size * sizeof(float),
+        input2_dims,
+        input2_dim_count,
+        ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
+        &input_tensor2
+    );
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        g_ort->ReleaseValue(input_tensor1);
+        return nullptr;
+    }
+
+    // 准备输入
+    const char* input_names[] = {sess->input_name_ptrs[0], sess->input_name_ptrs[1]};
+    OrtValue* input_tensors[] = {input_tensor1, input_tensor2};
+
+    // 运行推理
+    OrtValue* output_tensor = nullptr;
+    status = g_ort->Run(
+        sess->session,
+        nullptr,
+        input_names,
+        input_tensors,
+        2,
+        sess->output_name_ptrs.data(),
+        1,
+        &output_tensor
+    );
+
+    g_ort->ReleaseValue(input_tensor1);
+    g_ort->ReleaseValue(input_tensor2);
+
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        return nullptr;
+    }
+
+    // 获取输出数据
+    float* output_data = nullptr;
+    status = g_ort->GetTensorMutableData(output_tensor, (void**)&output_data);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        g_ort->ReleaseValue(output_tensor);
+        return nullptr;
+    }
+
+    // 获取输出大小
+    OrtTensorTypeAndShapeInfo* type_info = nullptr;
+    status = g_ort->GetTensorTypeAndShape(output_tensor, &type_info);
+    if (status) {
+        last_error = g_ort->GetErrorMessage(status);
+        g_ort->ReleaseStatus(status);
+        g_ort->ReleaseValue(output_tensor);
+        return nullptr;
+    }
+
+    size_t element_count = 0;
+    g_ort->GetTensorShapeElementCount(type_info, &element_count);
+    *output_size = (int)element_count;
+
+    // 复制输出数据
+    float* result = (float*)malloc(element_count * sizeof(float));
+    memcpy(result, output_data, element_count * sizeof(float));
+
+    g_ort->ReleaseTensorTypeAndShapeInfo(type_info);
+    g_ort->ReleaseValue(output_tensor);
+
+    return result;
+}