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

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,6 +58,7 @@ RUN mkdir -p /app/data /app/models
 
 ENV GIN_MODE=release
 ENV TZ=Asia/Shanghai
+ENV LD_LIBRARY_PATH=/usr/lib
 
 EXPOSE 6688
 

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

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=

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))
+}

pkg/onnx/onnx.cpp

@@ -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;
+}

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 <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,26 +57,124 @@ func LoadModel(name, path string) error {
 	return nil
 }
 
-// Run 执行模型推理
-func (s *Session) Run(input []float32, inputSize int) ([]float32, error) {
+// GetSession 获取已加载的会话
+func GetSession(name string) (*Session, bool) {
+	mu.RLock()
+	s, ok := sessions[name]
+	mu.RUnlock()
+	return s, ok
+}
+
+// Run 执行单输入推理
+func (s *Session) Run(input []float32, dims []int64) ([]float32, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	output := make([]float32, inputSize)
+	if len(input) == 0 {
+		return nil, errors.New("输入数据为空")
+	}
 
-	result := C.onnx_run(
+	var outputSize C.int
+
+	output := C.onnx_run_float(
 		s.session,
 		(*C.float)(unsafe.Pointer(&input[0])),
 		C.int(len(input)),
-		(*C.float)(unsafe.Pointer(&output[0])),
-		C.int(len(output)),
+		(*C.int64_t)(unsafe.Pointer(&dims[0])),
+		C.int(len(dims)),
+		&outputSize,
 	)
 
-	if result != 0 {
-		return nil, errors.New(C.GoString(C.onnx_get_last_error()))
+	if output == nil {
+		errMsg := C.GoString(C.onnx_get_last_error())
+		return nil, fmt.Errorf("推理失败: %s", errMsg)
 	}
 
-	return output, nil
+	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 关闭会话
@@ -99,14 +188,6 @@ func (s *Session) Close() {
 	}
 }
 
-// GetSession 获取已加载的会话
-func GetSession(name string) (*Session, bool) {
-	mu.RLock()
-	s, ok := sessions[name]
-	mu.RUnlock()
-	return s, ok
-}
-
 // CloseAll 关闭所有会话
 func CloseAll() {
 	mu.Lock()
@@ -116,3 +197,7 @@ func CloseAll() {
 	sessions = make(map[string]*Session)
 	mu.Unlock()
 }
+
+func init() {
+	C.onnx_init()
+}