diff --git a/Dockerfile b/Dockerfile index cb02bce..3dd38d2 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,22 +1,35 @@ # 构建阶段 - 使用 Debian 基础镜像 FROM golang:1.22-bookworm AS builder -# 安装 OpenCV 和编译工具 +# 安装 OpenCV、ONNX Runtime 和编译工具 RUN apt-get update && apt-get install -y --no-install-recommends \ gcc \ g++ \ libopencv-dev \ pkg-config \ + wget \ && rm -rf /var/lib/apt/lists/* +# 安装 ONNX Runtime +RUN wget -q https://github.com/microsoft/onnxruntime/releases/download/v1.16.3/onnxruntime-linux-x64-1.16.3.tgz \ + && tar -xzf onnxruntime-linux-x64-1.16.3.tgz \ + && mv onnxruntime-linux-x64-1.16.3 /usr/local/onnxruntime \ + && rm onnxruntime-linux-x64-1.16.3.tgz + +ENV ONNXRUNTIME_DIR=/usr/local/onnxruntime +ENV PKG_CONFIG_PATH=/usr/lib/x86_64-linux-gnu/pkgconfig +ENV CGO_ENABLED=1 + WORKDIR /app # 复制源代码 COPY . . -# 使用 pkg-config 获取 OpenCV 编译标志 -RUN go mod download || true -RUN PKG_CONFIG_PATH=/usr/lib/x86_64-linux-gnu/pkgconfig CGO_ENABLED=1 go build -ldflags="-s -w" -o anticaptcha ./cmd/server +# 下载依赖 +RUN go mod download + +# 构建 +RUN go build -ldflags="-s -w" -o anticaptcha ./cmd/server # 运行阶段 FROM debian:bookworm-slim @@ -31,6 +44,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ libstdc++6 \ && rm -rf /var/lib/apt/lists/* +# 复制 ONNX Runtime +COPY --from=builder /usr/local/onnxruntime/lib/libonnxruntime.so* /usr/lib/ + WORKDIR /app # 复制二进制文件和前端 @@ -42,7 +58,8 @@ RUN mkdir -p /app/data /app/models ENV GIN_MODE=release ENV TZ=Asia/Shanghai +ENV LD_LIBRARY_PATH=/usr/lib EXPOSE 6688 -CMD ["./anticaptcha"] +CMD ["./anticaptcha"] \ No newline at end of file diff --git a/go.mod b/go.mod index b0f480c..6b49a5b 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module anticaptcha go 1.22 require ( + github.com/disintegration/imaging v1.6.2 github.com/gin-contrib/cors v1.7.2 github.com/gin-gonic/gin v1.10.0 github.com/golang-jwt/jwt/v5 v5.2.1 @@ -52,6 +53,7 @@ require ( golang.org/x/arch v0.8.0 // indirect golang.org/x/crypto v0.25.0 // indirect golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect + golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8 // indirect golang.org/x/net v0.25.0 // indirect golang.org/x/sys v0.22.0 // indirect golang.org/x/text v0.16.0 // indirect diff --git a/go.sum b/go.sum index 64e5b80..8bb002a 100644 --- a/go.sum +++ b/go.sum @@ -10,6 +10,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/disintegration/imaging v1.6.2 h1:w1LecBlG2Lnp8B3jk5zSuNqd7b4DXhcjwek1ei82L+c= +github.com/disintegration/imaging v1.6.2/go.mod h1:44/5580QXChDfwIclfc/PCwrr44amcmDAg8hxG0Ewe4= github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA= @@ -120,12 +122,15 @@ golang.org/x/crypto v0.25.0 h1:ypSNr+bnYL2YhwoMt2zPxHFmbAN1KZs/njMG3hxUp30= golang.org/x/crypto v0.25.0/go.mod h1:T+wALwcMOSE0kXgUAnPAHqTLW+XHgcELELW8VaDgm/M= golang.org/x/exp v0.0.0-20230905200255-921286631fa9 h1:GoHiUyI/Tp2nVkLI2mCxVkOjsbSXD66ic0XW0js0R9g= golang.org/x/exp v0.0.0-20230905200255-921286631fa9/go.mod h1:S2oDrQGGwySpoQPVqRShND87VCbxmc6bL1Yd2oYrm6k= +golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8 h1:hVwzHzIUGRjiF7EcUjqNxk3NCfkPxbDKRdnNE1Rpg0U= +golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac= golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI= golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4= golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI= google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg= diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go index cce2a12..96a9a72 100644 --- a/internal/captcha/handler.go +++ b/internal/captcha/handler.go @@ -1,13 +1,22 @@ package captcha import ( + "bufio" + "bytes" + "encoding/base64" "fmt" + "image" + "math" "net/http" "os" "path/filepath" + "strings" "sync" + "anticaptcha/pkg/onnx" "anticaptcha/pkg/opencv" + + "github.com/disintegration/imaging" ) type Handler struct { @@ -33,9 +42,54 @@ func NewHandler(modelPath string) *Handler { } // 确保模型目录存在并下载缺失的模型 h.ensureModels() + // 加载模型 + h.loadModels() return h } +// loadModels 加载所有模型 +func (h *Handler) loadModels() { + // 加载 OCR 模型 + ocrPath := filepath.Join(h.modelPath, "OCR.onnx") + if _, err := os.Stat(ocrPath); err == nil { + if err := onnx.LoadModel("ocr", ocrPath); err != nil { + fmt.Printf("警告: 加载 OCR 模型失败: %v\n", err) + } else { + fmt.Println("OCR 模型加载成功") + } + } + + // 加载 Math 模型 + mathPath := filepath.Join(h.modelPath, "CRNN_Math.onnx") + if _, err := os.Stat(mathPath); err == nil { + if err := onnx.LoadModel("math", mathPath); err != nil { + fmt.Printf("警告: 加载 Math 模型失败: %v\n", err) + } else { + fmt.Println("Math 模型加载成功") + } + } + + // 加载 Rotation 模型 + rotatePath := filepath.Join(h.modelPath, "Rotation-RotNetR.onnx") + if _, err := os.Stat(rotatePath); err == nil { + if err := onnx.LoadModel("rotate", rotatePath); err != nil { + fmt.Printf("警告: 加载 Rotation 模型失败: %v\n", err) + } else { + fmt.Println("Rotation 模型加载成功") + } + } + + // 加载 Siamese 模型 + siamesePath := filepath.Join(h.modelPath, "Siamese-ResNet18.onnx") + if _, err := os.Stat(siamesePath); err == nil { + if err := onnx.LoadModel("siamese", siamesePath); err != nil { + fmt.Printf("警告: 加载 Siamese 模型失败: %v\n", err) + } else { + fmt.Println("Siamese 模型加载成功") + } + } +} + // ensureModels 检查并下载缺失的模型 func (h *Handler) ensureModels() { // 确保目录存在 @@ -81,20 +135,285 @@ func (h *Handler) downloadModel(remoteName, localPath string) error { return err } -// OCR 文字识别（需要 ONNX 模型） +// ===================== OCR 文字识别 ===================== + func (h *Handler) OCR(imageBase64 string) (string, error) { - // 暂时返回模拟结果 - // 实际实现需要加载 OCR 模型 - return "OCR result", nil + sess, ok := onnx.GetSession("ocr") + if !ok { + return "", fmt.Errorf("OCR 模型未加载") + } + + // 解码图片 + img, err := decodeBase64ToImage(imageBase64) + if err != nil { + return "", err + } + + // 加载字符集 + charset, err := h.loadCharset() + if err != nil { + return "", err + } + + // 预处理 + input, width, err := preprocessOCR(img) + if err != nil { + return "", err + } + + // 推理 + dims := []int64{1, 1, 64, int64(width)} + output, err := sess.Run(input, dims) + if err != nil { + return "", err + } + + // CTC 解码 + return ctcDecode(output, charset), nil } -// Math 数学计算识别 +// preprocessOCR OCR 预处理 +func preprocessOCR(img image.Image) ([]float32, int, error) { + // 调整高度为 64，保持宽高比 + bounds := img.Bounds() + width := bounds.Dx() + height := bounds.Dy() + newHeight := 64 + newWidth := width * newHeight / height + + resized := imaging.Resize(img, newWidth, newHeight, imaging.Lanczos) + gray := imaging.Grayscale(resized) + + // 转换为模型输入 + pixels := make([]float32, newWidth*newHeight) + for y := 0; y < newHeight; y++ { + for x := 0; x < newWidth; x++ { + c := gray.At(x, y) + r, _, _, _ := c.RGBA() + val := float32(r) / 65535.0 + pixels[y*newWidth+x] = (val - 0.5) / 0.5 + } + } + + return pixels, newWidth, nil +} + +// ctcDecode CTC 解码 +func ctcDecode(output []float32, charset []string) string { + if len(charset) == 0 { + return "" + } + + result := "" + lastIdx := 0 + numClasses := len(charset) + timesteps := len(output) / numClasses + + for t := 0; t < timesteps; t++ { + maxIdx := 0 + maxProb := float32(-math.MaxFloat32) + + for c := 0; c < numClasses; c++ { + idx := t * numClasses + c + if idx < len(output) && output[idx] > maxProb { + maxProb = output[idx] + maxIdx = c + } + } + + if maxIdx != 0 && maxIdx != lastIdx && maxIdx < len(charset) { + result += charset[maxIdx] + } + lastIdx = maxIdx + } + + return result +} + +// loadCharset 加载字符集 +func (h *Handler) loadCharset() ([]string, error) { + charsetPath := filepath.Join(h.modelPath, "CharSets.txt") + file, err := os.Open(charsetPath) + if err != nil { + return nil, err + } + defer file.Close() + + charset := make([]string, 0, 6000) + scanner := bufio.NewScanner(file) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line != "" { + charset = append(charset, line) + } + } + + // 添加空白符作为第一个字符 + result := make([]string, len(charset)+1) + result[0] = "" + copy(result[1:], charset) + + return result, scanner.Err() +} + +// ===================== Math 数学计算 ===================== + +const mathChars = "0123456789+-*/÷×=?" + func (h *Handler) Math(imageBase64 string) (string, error) { - // 暂时返回模拟结果 - return "0", nil + sess, ok := onnx.GetSession("math") + if !ok { + return "", fmt.Errorf("Math 模型未加载") + } + + // 解码图片 + img, err := decodeBase64ToImage(imageBase64) + if err != nil { + return "", err + } + + // 预处理 + input, err := preprocessMath(img) + if err != nil { + return "", err + } + + // 推理 + dims := []int64{1, 3, 70, 200} + output, err := sess.Run(input, dims) + if err != nil { + return "", err + } + + // 解码表达式 + expr := decodeMath(output) + if expr == "" { + return "", fmt.Errorf("无法识别表达式") + } + + // 计算结果 + result, err := evalMathExpression(expr) + if err != nil { + return "", err + } + + return fmt.Sprintf("%v", result), nil } -// SliderMatch 滑块缺口匹配 +// preprocessMath Math 预处理 +func preprocessMath(img image.Image) ([]float32, error) { + // 调整大小为 200x70，保持比例 + resized := imaging.Resize(img, 200, 70, imaging.Lanczos) + + // 转换为 RGB + rgb := imaging.Clone(resized) + + // 归一化 [N, C, H, W] + pixels := make([]float32, 3*70*200) + for y := 0; y < 70; y++ { + for x := 0; x < 200; x++ { + c := rgb.At(x, y) + r, g, b, _ := c.RGBA() + // CHW 格式，归一化 + pixels[0*70*200+y*200+x] = (float32(r)/65535.0 - 0.5) / 0.5 + pixels[1*70*200+y*200+x] = (float32(g)/65535.0 - 0.5) / 0.5 + pixels[2*70*200+y*200+x] = (float32(b)/65535.0 - 0.5) / 0.5 + } + } + + return pixels, nil +} + +// decodeMath 解码数学表达式 +func decodeMath(output []float32) string { + numChars := len(mathChars) + 1 + timesteps := len(output) / numChars + + result := "" + lastIdx := 0 + + for t := 0; t < timesteps; t++ { + maxIdx := 0 + maxProb := float32(-math.MaxFloat32) + + for c := 0; c < numChars; c++ { + idx := t * numChars + c + if idx < len(output) && output[idx] > maxProb { + maxProb = output[idx] + maxIdx = c + } + } + + if maxIdx != 0 && maxIdx != lastIdx { + if maxIdx-1 < len(mathChars) { + result += string(mathChars[maxIdx-1]) + } + } + lastIdx = maxIdx + } + + return result +} + +// evalMathExpression 计算数学表达式 +func evalMathExpression(expr string) (interface{}, error) { + // 替换特殊符号 + expr = strings.ReplaceAll(expr, "×", "*") + expr = strings.ReplaceAll(expr, "÷", "/") + expr = strings.ReplaceAll(expr, "?", "") + expr = strings.ReplaceAll(expr, "=", "") + + // 简单计算 + // 注意：实际项目中应使用更安全的方式 + var result float64 + var op byte = '+' + num := 0.0 + + for i := 0; i < len(expr); i++ { + c := expr[i] + if c >= '0' && c <= '9' { + num = num*10 + float64(c-'0') + } else if c == '+' || c == '-' || c == '*' || c == '/' { + switch op { + case '+': + result += num + case '-': + result -= num + case '*': + result *= num + case '/': + if num != 0 { + result /= num + } + } + op = c + num = 0 + } + } + + // 处理最后一个数字 + switch op { + case '+': + result += num + case '-': + result -= num + case '*': + result *= num + case '/': + if num != 0 { + result /= num + } + } + + // 返回整数或浮点数 + if result == float64(int(result)) { + return int(result), nil + } + return result, nil +} + +// ===================== 滑块匹配 ===================== + func (h *Handler) SliderMatch(targetBase64, backgroundBase64 string) (int, error) { target, err := opencv.DecodeFromBase64(targetBase64) if err != nil { @@ -111,7 +430,6 @@ func (h *Handler) SliderMatch(targetBase64, backgroundBase64 string) (int, error return opencv.SliderMatch(target, background) } -// SliderComparison 阴影滑块匹配 func (h *Handler) SliderComparison(targetBase64, backgroundBase64 string) (int, error) { target, err := opencv.DecodeFromBase64(targetBase64) if err != nil { @@ -128,8 +446,16 @@ func (h *Handler) SliderComparison(targetBase64, backgroundBase64 string) (int, return opencv.SliderComparison(target, background) } -// CompareSimilarity 图片相似度对比 +// ===================== 图像相似度 ===================== + func (h *Handler) CompareSimilarity(img1Base64, img2Base64 string) (float32, error) { + // 使用 ONNX 模型 + sess, ok := onnx.GetSession("siamese") + if ok { + return h.compareSimilarityONNX(sess, img1Base64, img2Base64) + } + + // 使用 OpenCV 直方图比较 img1, err := opencv.DecodeFromBase64(img1Base64) if err != nil { return 0, err @@ -145,8 +471,89 @@ func (h *Handler) CompareSimilarity(img1Base64, img2Base64 string) (float32, err return opencv.CompareSimilarity(img1, img2) } -// SingleRotate 单图旋转验证码 +func (h *Handler) compareSimilarityONNX(sess *onnx.Session, img1Base64, img2Base64 string) (float32, error) { + img1, err := decodeBase64ToImage(img1Base64) + if err != nil { + return 0, err + } + + img2, err := decodeBase64ToImage(img2Base64) + if err != nil { + return 0, err + } + + // 预处理 + input1, err := preprocessSiamese(img1) + if err != nil { + return 0, err + } + + input2, err := preprocessSiamese(img2) + if err != nil { + return 0, err + } + + // 推理 + dims := []int64{1, 3, 105, 105} + output, err := sess.RunDualInput(input1, dims, input2, dims) + if err != nil { + return 0, err + } + + // 计算相似度 + if len(output) >= 2 { + emb1 := output[:len(output)/2] + emb2 := output[len(output)/2:] + + // 欧氏距离 + var dist float32 + for i := 0; i < len(emb1); i++ { + d := emb1[i] - emb2[i] + dist += d * d + } + dist = float32(math.Sqrt(float64(dist))) + + // 相似度 + similarity := 1.0 / (1.0 + dist) + return similarity, nil + } + + return 0, fmt.Errorf("输出格式错误") +} + +func preprocessSiamese(img image.Image) ([]float32, error) { + // 调整大小为 105x105 + resized := imaging.Resize(img, 105, 105, imaging.Lanczos) + rgb := imaging.Clone(resized) + + // ImageNet 归一化 + mean := [3]float32{0.485, 0.456, 0.406} + std := [3]float32{0.229, 0.224, 0.225} + + pixels := make([]float32, 3*105*105) + for y := 0; y < 105; y++ { + for x := 0; x < 105; x++ { + c := rgb.At(x, y) + r, g, b, _ := c.RGBA() + pixels[0*105*105+y*105+x] = (float32(r)/65535.0 - mean[0]) / std[0] + pixels[1*105*105+y*105+x] = (float32(g)/65535.0 - mean[1]) / std[1] + pixels[2*105*105+y*105+x] = (float32(b)/65535.0 - mean[2]) / std[2] + } + } + + return pixels, nil +} + +// ===================== 旋转检测 ===================== + func (h *Handler) SingleRotate(imageBase64 string) (float32, error) { + // 使用 ONNX 模型 + sess, ok := onnx.GetSession("rotate") + if ok { + return h.singleRotateONNX(sess, imageBase64) + } + + // 使用 OpenCV img, err := opencv.DecodeFromBase64(imageBase64) if err != nil { return 0, err @@ -156,9 +563,62 @@ func (h *Handler) SingleRotate(imageBase64 string) (float32, error) { return opencv.DetectRotation(img) } -// DoubleRotate 双图旋转验证码 +func (h *Handler) singleRotateONNX(sess *onnx.Session, imageBase64 string) (float32, error) { + img, err := decodeBase64ToImage(imageBase64) + if err != nil { + return 0, err + } + + // 预处理 + input, err := preprocessRotation(img) + if err != nil { + return 0, err + } + + // 推理 + dims := []int64{1, 3, 224, 224} + output, err := sess.Run(input, dims) + if err != nil { + return 0, err + } + + // 找到最大概率的角度 + maxIdx := 0 + maxProb := float32(-math.MaxFloat32) + for i := 0; i < len(output); i++ { + if output[i] > maxProb { + maxProb = output[i] + maxIdx = i + } + } + + return float32(maxIdx), nil +} + +func preprocessRotation(img image.Image) ([]float32, error) { + // 调整大小为 224x224 + resized := imaging.Resize(img, 224, 224, imaging.Lanczos) + rgb := imaging.Clone(resized) + + // ImageNet 归一化 + mean := [3]float32{0.485, 0.456, 0.406} + std := [3]float32{0.229, 0.224, 0.225} + + pixels := make([]float32, 3*224*224) + for y := 0; y < 224; y++ { + for x := 0; x < 224; x++ { + c := rgb.At(x, y) + r, g, b, _ := c.RGBA() + pixels[0*224*224+y*224+x] = (float32(r)/65535.0 - mean[0]) / std[0] + pixels[1*224*224+y*224+x] = (float32(g)/65535.0 - mean[1]) / std[1] + pixels[2*224*224+y*224+x] = (float32(b)/65535.0 - mean[2]) / std[2] + } + } + + return pixels, nil +} + func (h *Handler) DoubleRotate(insideBase64, outsideBase64 string) (float32, error) { - // 简化处理 inside, err := opencv.DecodeFromBase64(insideBase64) if err != nil { return 0, err @@ -184,27 +644,43 @@ func (h *Handler) DoubleRotate(insideBase64, outsideBase64 string) (float32, err return angleInside - angleOutside, nil } -// DetectionIcon 图标检测 +// ===================== 图标/文字检测 ===================== + func (h *Handler) DetectionIcon(imageBase64 string) ([]map[string]int, error) { // 暂时返回空结果 - // 实际需要目标检测模型 return []map[string]int{}, nil } -// DetectionText 文字检测 func (h *Handler) DetectionText(imageBase64 string) ([]map[string]int, error) { - // 暂时返回空结果 return []map[string]int{}, nil } -// DetectionIconOrder 按序检测图标 func (h *Handler) DetectionIconOrder(orderImgBase64, targetImgBase64 string) ([]map[string]int, error) { - // 暂时返回空结果 return []map[string]int{}, nil } -// DetectionTextOrder 按序检测文字 func (h *Handler) DetectionTextOrder(orderImgBase64, targetImgBase64 string) ([]map[string]int, error) { - // 暂时返回空结果 return []map[string]int{}, nil } + +// ===================== 工具函数 ===================== + +func decodeBase64ToImage(base64Str string) (image.Image, error) { + data, err := base64.StdEncoding.DecodeString(base64Str) + if err != nil { + // 尝试去掉 data URL 前缀 + if strings.Contains(base64Str, ",") { + parts := strings.SplitN(base64Str, ",", 2) + if len(parts) == 2 { + data, err = base64.StdEncoding.DecodeString(parts[1]) + if err != nil { + return nil, err + } + } + } else { + return nil, err + } + } + + return imaging.Decode(bytes.NewReader(data)) +} \ No newline at end of file diff --git a/pkg/onnx/onnx.cpp b/pkg/onnx/onnx.cpp new file mode 100644 index 0000000..fe9620f --- /dev/null +++ b/pkg/onnx/onnx.cpp @@ -0,0 +1,424 @@ +#include +#include +#include +#include +#include + +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 input_names; + std::vector output_names; + std::vector input_name_ptrs; + std::vector 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; +} \ No newline at end of file diff --git a/pkg/onnx/onnx.go b/pkg/onnx/onnx.go index 657a44b..e2c3bfd 100644 --- a/pkg/onnx/onnx.go +++ b/pkg/onnx/onnx.go @@ -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 -#include -// 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() } \ No newline at end of file