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
@@ -1,35 +1,24 @@
 package onnx

 /*
-#cgo CXXFLAGS: -std=c++17
-#cgo linux LDFLAGS: -lonnxruntime -ldl
-#cgo darwin LDFLAGS: -lonnxruntime -framework CoreFoundation
+#cgo CXXFLAGS: -std=c++17 -I/usr/include/onnxruntime
+#cgo linux LDFLAGS: -lonnxruntime -ldl -lstdc++
+#cgo darwin LDFLAGS: -lonnxruntime -framework CoreFoundation -lstdc++

 #include <stdlib.h>
-#include <string.h>

-// ONNX Runtime C API 声明
-#ifdef __cplusplus
-extern "C" {
-#endif
+typedef struct OnnxSession OnnxSession;

-typedef void* OrtSession;
-typedef void* OrtMemoryInfo;
-typedef void* OrtValue;
-
-// 初始化 ONNX 会话
-OrtSession onnx_create_session(const char* model_path);
-void onnx_destroy_session(OrtSession session);
-
-// 运行推理
-int onnx_run(OrtSession session, const float* input_data, int input_size, float* output_data, int output_size);
-
-// 错误信息
+// ONNX Runtime 函数
+int onnx_init();
+OnnxSession* onnx_create_session(const char* model_path);
+void onnx_destroy_session(OnnxSession* session);
+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);
+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);
+int onnx_get_input_shape(OnnxSession* sess, int input_index, int64_t* dims, int max_dims);
+int onnx_get_output_shape(OnnxSession* sess, int output_index, int64_t* dims, int max_dims);
+void onnx_free(void* ptr);
 const char* onnx_get_last_error();
-
-#ifdef __cplusplus
-}
-#endif
 */
 import "C"
 import (
@@ -39,8 +28,9 @@ import (
 	"unsafe"
 )

+// Session ONNX 会话
 type Session struct {
-	session C.OrtSession
+	session *C.OnnxSession
 	mu      sync.Mutex
 }

@@ -56,7 +46,8 @@ func LoadModel(name, path string) error {

 	session := C.onnx_create_session(cPath)
 	if session == nil {
-		return fmt.Errorf("加载模型失败: %s", C.GoString(C.onnx_get_last_error()))
+		errMsg := C.GoString(C.onnx_get_last_error())
+		return fmt.Errorf("加载模型失败: %s", errMsg)
 	}

 	mu.Lock()
@@ -66,39 +57,6 @@ func LoadModel(name, path string) error {
 	return nil
 }

-// Run 执行模型推理
-func (s *Session) Run(input []float32, inputSize int) ([]float32, error) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	output := make([]float32, inputSize)
-	
-	result := C.onnx_run(
-		s.session,
-		(*C.float)(unsafe.Pointer(&input[0])),
-		C.int(len(input)),
-		(*C.float)(unsafe.Pointer(&output[0])),
-		C.int(len(output)),
-	)
-
-	if result != 0 {
-		return nil, errors.New(C.GoString(C.onnx_get_last_error()))
-	}
-
-	return output, nil
-}
-
-// Close 关闭会话
-func (s *Session) Close() {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	
-	if s.session != nil {
-		C.onnx_destroy_session(s.session)
-		s.session = nil
-	}
-}
-
 // GetSession 获取已加载的会话
 func GetSession(name string) (*Session, bool) {
 	mu.RLock()
@@ -107,6 +65,129 @@ func GetSession(name string) (*Session, bool) {
 	return s, ok
 }

+// Run 执行单输入推理
+func (s *Session) Run(input []float32, dims []int64) ([]float32, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if len(input) == 0 {
+		return nil, errors.New("输入数据为空")
+	}
+
+	var outputSize C.int
+
+	output := C.onnx_run_float(
+		s.session,
+		(*C.float)(unsafe.Pointer(&input[0])),
+		C.int(len(input)),
+		(*C.int64_t)(unsafe.Pointer(&dims[0])),
+		C.int(len(dims)),
+		&outputSize,
+	)
+
+	if output == nil {
+		errMsg := C.GoString(C.onnx_get_last_error())
+		return nil, fmt.Errorf("推理失败: %s", errMsg)
+	}
+
+	defer C.onnx_free(unsafe.Pointer(output))
+
+	// 复制输出数据
+	result := make([]float32, int(outputSize))
+	outputSlice := (*[1 << 30]float32)(unsafe.Pointer(output))[:int(outputSize):int(outputSize)]
+	for i := 0; i < int(outputSize); i++ {
+		result[i] = outputSlice[i]
+	}
+
+	return result, nil
+}
+
+// RunDualInput 执行双输入推理（用于 Siamese 网络）
+func (s *Session) RunDualInput(input1 []float32, dims1 []int64, input2 []float32, dims2 []int64) ([]float32, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if len(input1) == 0 || len(input2) == 0 {
+		return nil, errors.New("输入数据为空")
+	}
+
+	var outputSize C.int
+
+	output := C.onnx_run_dual_input(
+		s.session,
+		(*C.float)(unsafe.Pointer(&input1[0])),
+		C.int(len(input1)),
+		(*C.int64_t)(unsafe.Pointer(&dims1[0])),
+		C.int(len(dims1)),
+		(*C.float)(unsafe.Pointer(&input2[0])),
+		C.int(len(input2)),
+		(*C.int64_t)(unsafe.Pointer(&dims2[0])),
+		C.int(len(dims2)),
+		&outputSize,
+	)
+
+	if output == nil {
+		errMsg := C.GoString(C.onnx_get_last_error())
+		return nil, fmt.Errorf("推理失败: %s", errMsg)
+	}
+
+	defer C.onnx_free(unsafe.Pointer(output))
+
+	// 复制输出数据
+	result := make([]float32, int(outputSize))
+	outputSlice := (*[1 << 30]float32)(unsafe.Pointer(output))[:int(outputSize):int(outputSize)]
+	for i := 0; i < int(outputSize); i++ {
+		result[i] = outputSlice[i]
+	}
+
+	return result, nil
+}
+
+// GetInputShape 获取输入形状
+func (s *Session) GetInputShape(index int) ([]int64, error) {
+	var dims [8]C.int64_t
+	dimCount := C.onnx_get_input_shape(s.session, C.int(index), &dims[0], 8)
+
+	if dimCount < 0 {
+		return nil, fmt.Errorf("获取输入形状失败: %s", C.GoString(C.onnx_get_last_error()))
+	}
+
+	result := make([]int64, dimCount)
+	for i := 0; i < int(dimCount); i++ {
+		result[i] = int64(dims[i])
+	}
+
+	return result, nil
+}
+
+// GetOutputShape 获取输出形状
+func (s *Session) GetOutputShape(index int) ([]int64, error) {
+	var dims [8]C.int64_t
+	dimCount := C.onnx_get_output_shape(s.session, C.int(index), &dims[0], 8)
+
+	if dimCount < 0 {
+		return nil, fmt.Errorf("获取输出形状失败: %s", C.GoString(C.onnx_get_last_error()))
+	}
+
+	result := make([]int64, dimCount)
+	for i := 0; i < int(dimCount); i++ {
+		result[i] = int64(dims[i])
+	}
+
+	return result, nil
+}
+
+// Close 关闭会话
+func (s *Session) Close() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.session != nil {
+		C.onnx_destroy_session(s.session)
+		s.session = nil
+	}
+}
+
 // CloseAll 关闭所有会话
 func CloseAll() {
 	mu.Lock()
@@ -115,4 +196,8 @@ func CloseAll() {
 	}
 	sessions = make(map[string]*Session)
 	mu.Unlock()
+}
+
+func init() {
+	C.onnx_init()
 }