admin
524c404194
Features: - Go + CGO ONNX/OpenCV wrapper for high performance - SQLite (default) / MySQL database support - Optional Redis caching - JWT authentication system - Multiple captcha recognition APIs: - OCR text recognition - Slider captcha matching - Image similarity comparison - Rotation captcha detection - Object detection - React frontend with install wizard - Docker and docker-compose support - Gitea CI/CD pipeline Project structure: - cmd/server: Main entry point - internal/: Core business logic - pkg/onnx: ONNX Runtime CGO wrapper - pkg/opencv: OpenCV CGO wrapper - web/: React frontend - deploy/: Deployment configs - scripts/: Utility scripts
118 lines
2.1 KiB
Go
118 lines
2.1 KiB
Go
package onnx
|
|
|
|
/*
|
|
#cgo CXXFLAGS: -std=c++17
|
|
#cgo linux LDFLAGS: -lonnxruntime -ldl
|
|
#cgo darwin LDFLAGS: -lonnxruntime -framework CoreFoundation
|
|
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
// ONNX Runtime C API 声明
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
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);
|
|
|
|
// 错误信息
|
|
const char* onnx_get_last_error();
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
*/
|
|
import "C"
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"unsafe"
|
|
)
|
|
|
|
type Session struct {
|
|
session C.OrtSession
|
|
mu sync.Mutex
|
|
}
|
|
|
|
var (
|
|
sessions = make(map[string]*Session)
|
|
mu sync.RWMutex
|
|
)
|
|
|
|
// LoadModel 加载 ONNX 模型
|
|
func LoadModel(name, path string) error {
|
|
cPath := C.CString(path)
|
|
defer C.free(unsafe.Pointer(cPath))
|
|
|
|
session := C.onnx_create_session(cPath)
|
|
if session == nil {
|
|
return fmt.Errorf("加载模型失败: %s", C.GoString(C.onnx_get_last_error()))
|
|
}
|
|
|
|
mu.Lock()
|
|
sessions[name] = &Session{session: session}
|
|
mu.Unlock()
|
|
|
|
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()
|
|
s, ok := sessions[name]
|
|
mu.RUnlock()
|
|
return s, ok
|
|
}
|
|
|
|
// CloseAll 关闭所有会话
|
|
func CloseAll() {
|
|
mu.Lock()
|
|
for _, s := range sessions {
|
|
s.Close()
|
|
}
|
|
sessions = make(map[string]*Session)
|
|
mu.Unlock()
|
|
} |