1317 lines
31 KiB
Go
1317 lines
31 KiB
Go
package captcha
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"image"
|
||
"image/color"
|
||
"math"
|
||
"net/http"
|
||
"os"
|
||
"path/filepath"
|
||
"sort"
|
||
"strings"
|
||
"sync"
|
||
|
||
"anticaptcha/pkg/onnx"
|
||
"anticaptcha/pkg/opencv"
|
||
|
||
"github.com/disintegration/imaging"
|
||
)
|
||
|
||
type Handler struct {
|
||
modelPath string
|
||
mu sync.RWMutex
|
||
}
|
||
|
||
// 模型配置：本地文件名 -> 远程文件名（可选模型用 optional 标记）
|
||
var modelConfigs = map[string]struct {
|
||
remote string
|
||
optional bool
|
||
}{
|
||
"CRNN_Math.onnx": {remote: "[AntiCAP]-CRNN_Math.onnx", optional: false},
|
||
"OCR.onnx": {remote: "[Dddd]-OCR.onnx", optional: false},
|
||
"Rotation-RotNetR.onnx": {remote: "[AntiCAP]-Rotation-RotNetR.onnx", optional: false},
|
||
"Siamese-ResNet18.onnx": {remote: "[AntiCAP]-Siamese-ResNet18.onnx", optional: false},
|
||
"CharSets.txt": {remote: "[Dddd]-CharSets.txt", optional: false},
|
||
// Detection 模型需要 YOLO11 后处理支持，OpenCV 4.6 不兼容，暂标记为可选
|
||
// 后续可用 ONNX Runtime 实现或升级 OpenCV 4.9+
|
||
"Detection_Icon.onnx": {remote: "Detection_Icon.onnx", optional: true},
|
||
"Detection_Text.onnx": {remote: "Detection_Text.onnx", optional: true},
|
||
}
|
||
|
||
// 从 Gitea 仓库下载（公开仓库，无需认证）
|
||
const modelBaseURL = "https://git.viaeon.com/admin/anticaptcha/raw/branch/main/models"
|
||
|
||
func NewHandler(modelPath string) *Handler {
|
||
h := &Handler{
|
||
modelPath: modelPath,
|
||
}
|
||
// 确保模型目录存在并下载缺失的模型
|
||
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 模型加载成功")
|
||
}
|
||
}
|
||
|
||
// 加载 YOLO 检测模型（使用 ONNX Runtime 而非 OpenCV DNN，因为 YOLO11 不兼容 OpenCV 4.6）
|
||
iconPath := filepath.Join(h.modelPath, "Detection_Icon.onnx")
|
||
if _, err := os.Stat(iconPath); err == nil {
|
||
if err := onnx.LoadModel("detection_icon", iconPath); err != nil {
|
||
fmt.Printf("提示: Icon 检测模型加载失败: %v\n", err)
|
||
} else {
|
||
fmt.Println("Icon 检测模型加载成功 (ONNX Runtime)")
|
||
}
|
||
}
|
||
|
||
textPath := filepath.Join(h.modelPath, "Detection_Text.onnx")
|
||
if _, err := os.Stat(textPath); err == nil {
|
||
if err := onnx.LoadModel("detection_text", textPath); err != nil {
|
||
fmt.Printf("提示: Text 检测模型加载失败: %v\n", err)
|
||
} else {
|
||
fmt.Println("Text 检测模型加载成功 (ONNX Runtime)")
|
||
}
|
||
}
|
||
}
|
||
|
||
// ensureModels 检查并下载缺失的模型
|
||
func (h *Handler) ensureModels() {
|
||
if err := os.MkdirAll(h.modelPath, 0755); err != nil {
|
||
fmt.Printf("警告: 创建模型目录失败: %v\n", err)
|
||
return
|
||
}
|
||
|
||
for localName, cfg := range modelConfigs {
|
||
localPath := filepath.Join(h.modelPath, localName)
|
||
if _, err := os.Stat(localPath); os.IsNotExist(err) {
|
||
fmt.Printf("下载模型: %s -> %s\n", cfg.remote, localName)
|
||
if err := h.downloadModel(cfg.remote, localPath); err != nil {
|
||
if cfg.optional {
|
||
fmt.Printf("提示: 可选模型 %s 下载失败，相关功能不可用: %v\n", localName, err)
|
||
} else {
|
||
fmt.Printf("警告: 下载模型 %s 失败: %v\n", cfg.remote, err)
|
||
}
|
||
} else {
|
||
fmt.Printf("模型下载完成: %s\n", localName)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// downloadModel 下载模型文件
|
||
func (h *Handler) downloadModel(remoteName, localPath string) error {
|
||
url := fmt.Sprintf("%s/%s", modelBaseURL, remoteName)
|
||
|
||
resp, err := http.Get(url)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
return fmt.Errorf("HTTP %d", resp.StatusCode)
|
||
}
|
||
|
||
out, err := os.Create(localPath)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer out.Close()
|
||
|
||
_, err = out.ReadFrom(resp.Body)
|
||
return err
|
||
}
|
||
|
||
// ===================== OCR 文字识别 =====================
|
||
|
||
func (h *Handler) OCR(imageBase64 string) (string, error) {
|
||
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
|
||
}
|
||
|
||
return ctcDecode(output, charset), nil
|
||
}
|
||
|
||
func preprocessOCR(img image.Image) ([]float32, int, error) {
|
||
bounds := img.Bounds()
|
||
width := bounds.Dx()
|
||
height := bounds.Dy()
|
||
newHeight := 64
|
||
newWidth := width * newHeight / height
|
||
if newWidth < 1 {
|
||
newWidth = 1
|
||
}
|
||
|
||
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
|
||
}
|
||
|
||
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
|
||
}
|
||
|
||
func (h *Handler) loadCharset() ([]string, error) {
|
||
charsetPath := filepath.Join(h.modelPath, "CharSets.txt")
|
||
|
||
// 读取整个文件
|
||
data, err := os.ReadFile(charsetPath)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 解析 JSON 数组格式
|
||
var charset []string
|
||
if err := json.Unmarshal(data, &charset); err != nil {
|
||
return nil, fmt.Errorf("解析字符集失败: %v", err)
|
||
}
|
||
|
||
// 确保第一个元素是空字符串（CTC blank）
|
||
if len(charset) == 0 || charset[0] != "" {
|
||
charset = append([]string{""}, charset...)
|
||
}
|
||
|
||
return charset, nil
|
||
}
|
||
|
||
// ===================== Math 数学计算 =====================
|
||
|
||
const mathChars = "0123456789+-*/÷×=?"
|
||
|
||
func (h *Handler) Math(imageBase64 string) (result string, err error) {
|
||
// 捕获 panic
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
err = fmt.Errorf("Math panic: %v", r)
|
||
}
|
||
}()
|
||
|
||
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
|
||
}
|
||
|
||
// DEBUG: 打印输出前10个值
|
||
debugStr := "ONNX output first 10: "
|
||
for i := 0; i < 10 && i < len(output); i++ {
|
||
debugStr += fmt.Sprintf("%.4f ", output[i])
|
||
}
|
||
fmt.Println(debugStr)
|
||
|
||
// DEBUG: 对第一个位置 argmax
|
||
maxIdx := 0
|
||
maxProb := float32(-math.MaxFloat32)
|
||
for c := 0; c < 19; c++ {
|
||
if output[c] > maxProb {
|
||
maxProb = output[c]
|
||
maxIdx = c
|
||
}
|
||
}
|
||
// 【注意】这里用 string 索引访问多字节字符可能有问题，仅用于调试
|
||
debugChars := []string{"", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "+", "-", "*", "/", "÷", "×", "=", "?"}
|
||
if maxIdx >= 0 && maxIdx < len(debugChars) {
|
||
fmt.Printf("DEBUG: Position 0 argmax = %d (char: %s)\n", maxIdx, debugChars[maxIdx])
|
||
}
|
||
|
||
if len(output) == 0 {
|
||
return "", fmt.Errorf("模型输出为空")
|
||
}
|
||
|
||
expr := decodeMathFromOutput(output)
|
||
fmt.Printf("DEBUG: Decoded expression = '%s'\n", expr)
|
||
if expr == "" {
|
||
return "", fmt.Errorf("无法识别表达式")
|
||
}
|
||
|
||
evalResult, err := evalMathExpression(expr)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
return fmt.Sprintf("%v", evalResult), nil
|
||
}
|
||
|
||
func preprocessMath(img image.Image) ([]float32, error) {
|
||
// resize_with_padding: 保持比例缩放，白色填充
|
||
targetW, targetH := 200, 70
|
||
|
||
// 计算缩放比例
|
||
bounds := img.Bounds()
|
||
srcW, srcH := bounds.Dx(), bounds.Dy()
|
||
if srcW <= 0 || srcH <= 0 {
|
||
return nil, fmt.Errorf("无效的图片尺寸")
|
||
}
|
||
|
||
ratio := minFloat(float64(targetW)/float64(srcW), float64(targetH)/float64(srcH))
|
||
|
||
newW := int(float64(srcW) * ratio)
|
||
newH := int(float64(srcH) * ratio)
|
||
|
||
// 确保至少 1 像素
|
||
if newW <= 0 {
|
||
newW = 1
|
||
}
|
||
if newH <= 0 {
|
||
newH = 1
|
||
}
|
||
|
||
// 缩放图片
|
||
resized := imaging.Resize(img, newW, newH, imaging.Lanczos)
|
||
|
||
// 创建白色背景
|
||
canvas := image.NewRGBA(image.Rect(0, 0, targetW, targetH))
|
||
white := color.RGBA{255, 255, 255, 255}
|
||
for x := 0; x < targetW; x++ {
|
||
for y := 0; y < targetH; y++ {
|
||
canvas.Set(x, y, white)
|
||
}
|
||
}
|
||
|
||
// 粘贴缩放后的图片
|
||
for x := 0; x < newW; x++ {
|
||
for y := 0; y < newH; y++ {
|
||
canvas.Set(x, y, resized.At(x, y))
|
||
}
|
||
}
|
||
|
||
pixels := make([]float32, 3*70*200)
|
||
for y := 0; y < 70; y++ {
|
||
for x := 0; x < 200; x++ {
|
||
c := canvas.At(x, y)
|
||
r, g, b, _ := c.RGBA()
|
||
// 归一化到 [0, 1]，然后标准化到 [-1, 1]
|
||
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
|
||
}
|
||
|
||
func decodeMathFromOutput(output []float32) string {
|
||
// Math 模型输出: [T, B, C] = [51, 1, 19]
|
||
// 51 个时间步，batch=1，19 个字符类别（blank + 18 chars）
|
||
// 原版 Python: transpose(1, 0, 2) -> [B, T, C]，然后 argmax(axis=2)
|
||
// 但实际上对于 [T, B, C] 行优先存储，argmax 在 C 维度不需要改变读取顺序
|
||
const T = 51
|
||
const B = 1
|
||
const C = 19
|
||
|
||
if len(output) != T*B*C {
|
||
total := len(output)
|
||
if total%C != 0 {
|
||
return ""
|
||
}
|
||
}
|
||
|
||
// 字符集：索引 1-18 对应字符
|
||
// 【关键修复】使用 string slice 而不是 string，避免 UTF-8 多字节字符的索引错位问题
|
||
// Go 的 string 是 UTF-8 编码，for range 迭代时 i 是字节位置而非字符位置
|
||
// 模型输出的索引是按字符位置（1-18），直接用 string[idx-1] 访问多字节字符会错位
|
||
// 例如：÷ (U+00F7, 2字节) 和 × (U+00D7, 2字节) 会导致后续字符索引跳跃
|
||
CHARS := []string{"", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "+", "-", "*", "/", "÷", "×", "=", "?"}
|
||
|
||
// DEBUG: 打印字符集
|
||
fmt.Printf("DEBUG CHARS slice (len=%d): ", len(CHARS))
|
||
for i, c := range CHARS {
|
||
fmt.Printf("[%d]=%s ", i, c)
|
||
}
|
||
fmt.Println()
|
||
|
||
// 先收集所有时间步的 argmax 结果，用于调试
|
||
allChars := make([]int, T)
|
||
for t := 0; t < T; t++ {
|
||
maxIdx := 0
|
||
maxProb := float32(-math.MaxFloat32)
|
||
for c := 0; c < C; c++ {
|
||
// [T, B, C] 行优先存储: output[t*B*C + b*C + c]
|
||
idx := t*B*C + c
|
||
if idx < len(output) && output[idx] > maxProb {
|
||
maxProb = output[idx]
|
||
maxIdx = c
|
||
}
|
||
}
|
||
allChars[t] = maxIdx
|
||
}
|
||
|
||
// 打印所有位置的识别结果（包括索引值）
|
||
debugStr := "DEBUG argmax: "
|
||
for _, idx := range allChars {
|
||
if idx >= 0 && idx < len(CHARS) {
|
||
debugStr += CHARS[idx]
|
||
} else {
|
||
debugStr += fmt.Sprintf("[%d]", idx) // 显示越界索引
|
||
}
|
||
}
|
||
fmt.Println(debugStr)
|
||
|
||
// 打印原始索引值（全部 51 个）
|
||
fmt.Printf("DEBUG raw indices (all %d): %v\n", len(allChars), allChars)
|
||
|
||
// CTC 解码: 跳过 blank (index 0) 和连续重复
|
||
result := ""
|
||
lastIdx := -1
|
||
for _, idx := range allChars {
|
||
if idx != 0 && idx != lastIdx {
|
||
if idx < len(CHARS) {
|
||
result += CHARS[idx]
|
||
}
|
||
}
|
||
lastIdx = idx
|
||
}
|
||
|
||
fmt.Printf("DEBUG CTC decoded: '%s'\n", result)
|
||
return result
|
||
}
|
||
|
||
func minFloat(a, b float64) float64 {
|
||
if a < b {
|
||
return a
|
||
}
|
||
return b
|
||
}
|
||
|
||
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
|
||
}
|
||
|
||
// ===================== 滑块匹配 =====================
|
||
|
||
type SliderMatchResult struct {
|
||
Target []int `json:"target"`
|
||
}
|
||
|
||
func (h *Handler) SliderMatch(targetBase64, backgroundBase64 string) (*SliderMatchResult, error) {
|
||
target, err := opencv.DecodeFromBase64(targetBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer target.Free()
|
||
|
||
background, err := opencv.DecodeFromBase64(backgroundBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer background.Free()
|
||
|
||
x, err := opencv.SliderMatch(target, background)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return &SliderMatchResult{
|
||
Target: []int{x, 0, x + target.Width(), target.Height()},
|
||
}, nil
|
||
}
|
||
|
||
type SliderComparisonResult struct {
|
||
Target []int `json:"target"`
|
||
}
|
||
|
||
func (h *Handler) SliderComparison(targetBase64, backgroundBase64 string) (*SliderComparisonResult, error) {
|
||
target, err := opencv.DecodeFromBase64(targetBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer target.Free()
|
||
|
||
background, err := opencv.DecodeFromBase64(backgroundBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer background.Free()
|
||
|
||
x, y, err := opencv.SliderComparison(target, background)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return &SliderComparisonResult{
|
||
Target: []int{x, y},
|
||
}, nil
|
||
}
|
||
|
||
// ===================== 图像相似度 =====================
|
||
|
||
func (h *Handler) CompareSimilarity(img1Base64, img2Base64 string) (float32, error) {
|
||
sess, ok := onnx.GetSession("siamese")
|
||
if ok {
|
||
return h.compareSimilarityONNX(sess, img1Base64, img2Base64)
|
||
}
|
||
|
||
img1, err := opencv.DecodeFromBase64(img1Base64)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
defer img1.Free()
|
||
|
||
img2, err := opencv.DecodeFromBase64(img2Base64)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
defer img2.Free()
|
||
|
||
return opencv.CompareSimilarity(img1, img2)
|
||
}
|
||
|
||
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) {
|
||
resized := imaging.Resize(img, 105, 105, imaging.Lanczos)
|
||
rgb := imaging.Clone(resized)
|
||
|
||
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) (int, error) {
|
||
sess, ok := onnx.GetSession("rotate")
|
||
if ok {
|
||
return h.singleRotateONNX(sess, imageBase64)
|
||
}
|
||
|
||
img, err := opencv.DecodeFromBase64(imageBase64)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
defer img.Free()
|
||
|
||
angle, err := opencv.DetectRotation(img)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
|
||
return int(angle), nil
|
||
}
|
||
|
||
func (h *Handler) singleRotateONNX(sess *onnx.Session, imageBase64 string) (int, 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 maxIdx, nil
|
||
}
|
||
|
||
func preprocessRotation(img image.Image) ([]float32, error) {
|
||
bounds := img.Bounds()
|
||
w, h := bounds.Dx(), bounds.Dy()
|
||
|
||
size := w
|
||
if h < w {
|
||
size = h
|
||
}
|
||
|
||
cropX := (w - size) / 2
|
||
cropY := (h - size) / 2
|
||
cropped := imaging.Crop(img, image.Rect(cropX, cropY, cropX+size, cropY+size))
|
||
|
||
sqrt2 := math.Sqrt(2.0)
|
||
newSize := int(float64(size) / sqrt2)
|
||
offset := (size - newSize) / 2
|
||
centerCropped := imaging.Crop(cropped, image.Rect(offset, offset, offset+newSize, offset+newSize))
|
||
|
||
resized := imaging.Resize(centerCropped, 224, 224, imaging.Lanczos)
|
||
rgb := imaging.Clone(resized)
|
||
|
||
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
|
||
}
|
||
|
||
type DoubleRotateResult struct {
|
||
Angle int `json:"angle"`
|
||
}
|
||
|
||
func (h *Handler) DoubleRotate(insideBase64, outsideBase64 string, checkPixel int, speedRatio float64, grayscale, anticlockwise bool, cutPixelValue int) (*DoubleRotateResult, error) {
|
||
sess, ok := onnx.GetSession("rotate")
|
||
if ok {
|
||
insideAngle, err := h.singleRotateONNX(sess, insideBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
outsideAngle, err := h.singleRotateONNX(sess, outsideBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
angle := insideAngle - outsideAngle
|
||
if anticlockwise {
|
||
angle = -angle
|
||
}
|
||
if angle < 0 {
|
||
angle += 360
|
||
}
|
||
|
||
return &DoubleRotateResult{Angle: angle}, nil
|
||
}
|
||
|
||
inside, err := opencv.DecodeFromBase64(insideBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer inside.Free()
|
||
|
||
outside, err := opencv.DecodeFromBase64(outsideBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer outside.Free()
|
||
|
||
angleInside, err := opencv.DetectRotation(inside)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
angleOutside, err := opencv.DetectRotation(outside)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
angle := int(angleInside - angleOutside)
|
||
if anticlockwise {
|
||
angle = -angle
|
||
}
|
||
if angle < 0 {
|
||
angle += 360
|
||
}
|
||
|
||
return &DoubleRotateResult{Angle: angle}, nil
|
||
}
|
||
|
||
// ===================== 图标/文字检测 (YOLO via OpenCV DNN) =====================
|
||
|
||
type Detection struct {
|
||
Class string `json:"class"`
|
||
Box []int `json:"box"`
|
||
}
|
||
|
||
type DetectionResult struct {
|
||
Detections []Detection `json:"detections"`
|
||
}
|
||
|
||
func (h *Handler) DetectionIcon(imageBase64 string) (*DetectionResult, error) {
|
||
return h.detectYOLO("icon", imageBase64)
|
||
}
|
||
|
||
func (h *Handler) DetectionText(imageBase64 string) (*DetectionResult, error) {
|
||
return h.detectYOLO("text", imageBase64)
|
||
}
|
||
|
||
func (h *Handler) detectYOLO(modelName, imageBase64 string) (*DetectionResult, error) {
|
||
// YOLO11 模型使用 ONNX Runtime（OpenCV 4.6 DNN 不兼容 YOLO11 输出格式）
|
||
sess, ok := onnx.GetSession("detection_" + modelName)
|
||
if !ok {
|
||
return nil, fmt.Errorf("YOLO 检测模型 '%s' 未加载（需要 YOLO11 ONNX 模型）", modelName)
|
||
}
|
||
|
||
img, err := decodeBase64ToImage(imageBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 获取原图尺寸
|
||
origW := img.Bounds().Dx()
|
||
origH := img.Bounds().Dy()
|
||
|
||
// 预处理：Resize 到 640x640，归一化
|
||
input, err := preprocessYOLO(img)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
dims := []int64{1, 3, 640, 640}
|
||
output, err := sess.Run(input, dims)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// DEBUG: 打印 ONNX 输出的前 20 个值
|
||
fmt.Printf("DEBUG YOLO output first 20 values: ")
|
||
for i := 0; i < 20 && i < len(output); i++ {
|
||
fmt.Printf("%.4f ", output[i])
|
||
}
|
||
fmt.Println()
|
||
|
||
// YOLO11 输出: [1, 5, 8400] 或 [1, 84, 8400]
|
||
// 需要转置 + NMS 后处理，坐标从 640x640 缩放到原图尺寸
|
||
// 对于小图片（order 提示图）使用更低的置信度阈值
|
||
confThresh := 0.25
|
||
if origW < 200 || origH < 100 {
|
||
confThresh = 0.1 // 小图使用更低阈值
|
||
}
|
||
detections := postprocessYOLO(output, origW, origH, confThresh, 0.45)
|
||
|
||
result := &DetectionResult{
|
||
Detections: make([]Detection, len(detections)),
|
||
}
|
||
|
||
for i, det := range detections {
|
||
result.Detections[i] = Detection{
|
||
Class: det.ClassName,
|
||
Box: []int{int(det.X1), int(det.Y1), int(det.X2), int(det.Y2)},
|
||
}
|
||
}
|
||
|
||
fmt.Printf("DEBUG YOLO '%s': detected %d objects, origSize=%dx%d, confThresh=%.2f\n", modelName, len(detections), origW, origH, confThresh)
|
||
|
||
return result, nil
|
||
}
|
||
|
||
// YOLODetection YOLO 检测结果
|
||
type YOLODetection struct {
|
||
X1, Y1, X2, Y2 float64
|
||
Confidence float64
|
||
ClassID int
|
||
ClassName string
|
||
}
|
||
|
||
// preprocessYOLO 预处理图片为 YOLO 输入
|
||
func preprocessYOLO(img image.Image) ([]float32, error) {
|
||
// Resize 到 640x640
|
||
resized := imaging.Resize(img, 640, 640, imaging.Lanczos)
|
||
|
||
pixels := make([]float32, 3*640*640)
|
||
for y := 0; y < 640; y++ {
|
||
for x := 0; x < 640; x++ {
|
||
c := resized.At(x, y)
|
||
r, g, b, _ := c.RGBA()
|
||
// 归一化到 [0, 1]
|
||
pixels[0*640*640+y*640+x] = float32(r) / 65535.0
|
||
pixels[1*640*640+y*640+x] = float32(g) / 65535.0
|
||
pixels[2*640*640+y*640+x] = float32(b) / 65535.0
|
||
}
|
||
}
|
||
|
||
return pixels, nil
|
||
}
|
||
|
||
// postprocessYOLO YOLO11 后处理
|
||
func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh float64) []YOLODetection {
|
||
// YOLO11 ONNX 输出格式: [1, 5, 8400] 或 [1, 84, 8400]
|
||
// 布局: [batch, features, boxes]
|
||
// 对于单类检测 [1, 5, 8400]:
|
||
// output[0:8400] = 所有 box 的 x (中心点)
|
||
// output[8400:16800] = 所有 box 的 y (中心点)
|
||
// output[16800:25200] = 所有 box 的 w
|
||
// output[25200:33600] = 所有 box 的 h
|
||
// output[33600:42000] = 所有 box 的 confidence
|
||
|
||
outputLen := len(output)
|
||
if outputLen == 0 {
|
||
return nil
|
||
}
|
||
|
||
numBoxes := 8400
|
||
numFeatures := 5
|
||
|
||
if outputLen != numFeatures*numBoxes {
|
||
fmt.Printf("DEBUG YOLO: unexpected output length %d (expected %d)\n", outputLen, numFeatures*numBoxes)
|
||
return nil
|
||
}
|
||
|
||
// 类别名称
|
||
classNames := []string{"icon"}
|
||
|
||
// 坐标缩放比例：从 640x640 到原图
|
||
scaleX := float64(origW) / 640.0
|
||
scaleY := float64(origH) / 640.0
|
||
|
||
detections := []YOLODetection{}
|
||
maxConf := 0.0
|
||
|
||
// 遍历所有检测框
|
||
for i := 0; i < numBoxes; i++ {
|
||
// [features, boxes] 布局
|
||
// output[feat_idx * numBoxes + box_idx]
|
||
confIdx := 4*numBoxes + i // 第 5 个特征（索引 4）
|
||
confidence := float64(output[confIdx])
|
||
|
||
if confidence > maxConf {
|
||
maxConf = confidence
|
||
}
|
||
|
||
if confidence < confThresh {
|
||
continue
|
||
}
|
||
|
||
// 获取边界框 (x, y, w, h)
|
||
cx := float64(output[0*numBoxes + i])
|
||
cy := float64(output[1*numBoxes + i])
|
||
w := float64(output[2*numBoxes + i])
|
||
h := float64(output[3*numBoxes + i])
|
||
|
||
// 缩放坐标到原图尺寸
|
||
cx *= scaleX
|
||
cy *= scaleY
|
||
w *= scaleX
|
||
h *= scaleY
|
||
|
||
// 转换为 x1, y1, x2, y2
|
||
x1 := cx - w/2
|
||
y1 := cy - h/2
|
||
x2 := cx + w/2
|
||
y2 := cy + h/2
|
||
|
||
className := "object"
|
||
if 0 < len(classNames) {
|
||
className = classNames[0]
|
||
}
|
||
|
||
detections = append(detections, YOLODetection{
|
||
X1: x1,
|
||
Y1: y1,
|
||
X2: x2,
|
||
Y2: y2,
|
||
Confidence: confidence,
|
||
ClassID: 0,
|
||
ClassName: className,
|
||
})
|
||
}
|
||
|
||
fmt.Printf("DEBUG YOLO postprocess: outputLen=%d, format=[1,%d,%d], maxConf=%.4f, found=%d\n", outputLen, numFeatures, numBoxes, maxConf, len(detections))
|
||
|
||
// NMS
|
||
return nms(detections, nmsThresh)
|
||
}
|
||
|
||
// nms 非极大值抑制
|
||
func nms(detections []YOLODetection, thresh float64) []YOLODetection {
|
||
if len(detections) == 0 {
|
||
return nil
|
||
}
|
||
|
||
// 按置信度排序
|
||
sort.Slice(detections, func(i, j int) bool {
|
||
return detections[i].Confidence > detections[j].Confidence
|
||
})
|
||
|
||
keep := make([]bool, len(detections))
|
||
for i := range keep {
|
||
keep[i] = true
|
||
}
|
||
|
||
for i := 0; i < len(detections); i++ {
|
||
if !keep[i] {
|
||
continue
|
||
}
|
||
for j := i + 1; j < len(detections); j++ {
|
||
if !keep[j] {
|
||
continue
|
||
}
|
||
|
||
// 计算 IoU
|
||
iou := computeIoU(detections[i], detections[j])
|
||
if iou > thresh {
|
||
keep[j] = false
|
||
}
|
||
}
|
||
}
|
||
|
||
result := []YOLODetection{}
|
||
for i, det := range detections {
|
||
if keep[i] {
|
||
result = append(result, det)
|
||
}
|
||
}
|
||
|
||
return result
|
||
}
|
||
|
||
// computeIoU 计算 IoU
|
||
func computeIoU(a, b YOLODetection) float64 {
|
||
x1 := max(a.X1, b.X1)
|
||
y1 := max(a.Y1, b.Y1)
|
||
x2 := min(a.X2, b.X2)
|
||
y2 := min(a.Y2, b.Y2)
|
||
|
||
if x2 <= x1 || y2 <= y1 {
|
||
return 0
|
||
}
|
||
|
||
inter := (x2 - x1) * (y2 - y1)
|
||
areaA := (a.X2 - a.X1) * (a.Y2 - a.Y1)
|
||
areaB := (b.X2 - b.X1) * (b.Y2 - b.Y1)
|
||
union := areaA + areaB - inter
|
||
|
||
return inter / union
|
||
}
|
||
|
||
func max(a, b float64) float64 {
|
||
if a > b {
|
||
return a
|
||
}
|
||
return b
|
||
}
|
||
|
||
func min(a, b float64) float64 {
|
||
if a < b {
|
||
return a
|
||
}
|
||
return b
|
||
}
|
||
|
||
// ===================== 按序点击 (匈牙利算法匹配) =====================
|
||
|
||
func (h *Handler) DetectionIconOrder(orderImgBase64, targetImgBase64 string) ([]Detection, error) {
|
||
return h.detectOrder("icon", orderImgBase64, targetImgBase64)
|
||
}
|
||
|
||
func (h *Handler) DetectionTextOrder(orderImgBase64, targetImgBase64 string) ([]Detection, error) {
|
||
return h.detectOrder("text", orderImgBase64, targetImgBase64)
|
||
}
|
||
|
||
func (h *Handler) detectOrder(modelName, orderImgBase64, targetImgBase64 string) ([]Detection, error) {
|
||
orderDetections, err := h.detectYOLO(modelName, orderImgBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
targetDetections, err := h.detectYOLO(modelName, targetImgBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
sort.Slice(orderDetections.Detections, func(i, j int) bool {
|
||
return orderDetections.Detections[i].Box[0] < orderDetections.Detections[j].Box[0]
|
||
})
|
||
|
||
return h.hungarianMatch(orderImgBase64, targetImgBase64, orderDetections.Detections, targetDetections.Detections)
|
||
}
|
||
|
||
func (h *Handler) hungarianMatch(orderImgBase64, targetImgBase64 string, orderBoxes, targetBoxes []Detection) ([]Detection, error) {
|
||
if len(orderBoxes) == 0 || len(targetBoxes) == 0 {
|
||
return make([]Detection, len(orderBoxes)), nil
|
||
}
|
||
|
||
orderImg, err := decodeBase64ToImage(orderImgBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
targetImg, err := decodeBase64ToImage(targetImgBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
numOrders := len(orderBoxes)
|
||
numTargets := len(targetBoxes)
|
||
|
||
costMatrix := make([][]float64, numOrders)
|
||
for i := range costMatrix {
|
||
costMatrix[i] = make([]float64, numTargets)
|
||
for j := range costMatrix[i] {
|
||
costMatrix[i][j] = 1.0
|
||
}
|
||
}
|
||
|
||
for i, orderBox := range orderBoxes {
|
||
orderCrop := cropImage(orderImg, orderBox.Box)
|
||
if orderCrop == nil {
|
||
continue
|
||
}
|
||
|
||
for j, targetBox := range targetBoxes {
|
||
targetCrop := cropImage(targetImg, targetBox.Box)
|
||
if targetCrop == nil {
|
||
continue
|
||
}
|
||
|
||
similarity, err := h.computeImageSimilarity(orderCrop, targetCrop)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
|
||
costMatrix[i][j] = 1.0 - float64(similarity)
|
||
}
|
||
}
|
||
|
||
assignments := hungarian(costMatrix)
|
||
|
||
result := make([]Detection, numOrders)
|
||
for i, j := range assignments {
|
||
if j >= 0 && j < len(targetBoxes) {
|
||
result[i] = targetBoxes[j]
|
||
}
|
||
}
|
||
|
||
return result, nil
|
||
}
|
||
|
||
func cropImage(img image.Image, box []int) image.Image {
|
||
if len(box) < 4 {
|
||
return nil
|
||
}
|
||
bounds := img.Bounds()
|
||
if box[0] < 0 || box[1] < 0 || box[2] > bounds.Dx() || box[3] > bounds.Dy() {
|
||
return nil
|
||
}
|
||
if box[2] <= box[0] || box[3] <= box[1] {
|
||
return nil
|
||
}
|
||
return imaging.Crop(img, image.Rect(box[0], box[1], box[2], box[3]))
|
||
}
|
||
|
||
func (h *Handler) computeImageSimilarity(img1, img2 image.Image) (float32, error) {
|
||
sess, ok := onnx.GetSession("siamese")
|
||
if ok {
|
||
buf1 := new(bytes.Buffer)
|
||
imaging.Encode(buf1, img1, imaging.PNG)
|
||
b64_1 := base64.StdEncoding.EncodeToString(buf1.Bytes())
|
||
|
||
buf2 := new(bytes.Buffer)
|
||
imaging.Encode(buf2, img2, imaging.PNG)
|
||
b64_2 := base64.StdEncoding.EncodeToString(buf2.Bytes())
|
||
|
||
return h.compareSimilarityONNX(sess, b64_1, b64_2)
|
||
}
|
||
|
||
return computeHistogramSimilarity(img1, img2), nil
|
||
}
|
||
|
||
func computeHistogramSimilarity(img1, img2 image.Image) float32 {
|
||
size := 64
|
||
resized1 := imaging.Resize(img1, size, size, imaging.Lanczos)
|
||
resized2 := imaging.Resize(img2, size, size, imaging.Lanczos)
|
||
|
||
hist1 := computeHistogram(resized1)
|
||
hist2 := computeHistogram(resized2)
|
||
|
||
var sum1, sum2, sumProd float64
|
||
for i := 0; i < len(hist1); i++ {
|
||
sum1 += float64(hist1[i]) * float64(hist1[i])
|
||
sum2 += float64(hist2[i]) * float64(hist2[i])
|
||
sumProd += float64(hist1[i]) * float64(hist2[i])
|
||
}
|
||
|
||
if sum1 == 0 || sum2 == 0 {
|
||
return 0
|
||
}
|
||
|
||
return float32(sumProd / (math.Sqrt(sum1) * math.Sqrt(sum2)))
|
||
}
|
||
|
||
func computeHistogram(img image.Image) []int {
|
||
hist := make([]int, 256)
|
||
bounds := img.Bounds()
|
||
for y := 0; y < bounds.Dy(); y++ {
|
||
for x := 0; x < bounds.Dx(); x++ {
|
||
c := img.At(x, y)
|
||
r, g, b, _ := c.RGBA()
|
||
gray := int((r + g + b) / 3 / 256)
|
||
hist[gray]++
|
||
}
|
||
}
|
||
return hist
|
||
}
|
||
|
||
func hungarian(costMatrix [][]float64) []int {
|
||
n := len(costMatrix)
|
||
if n == 0 {
|
||
return nil
|
||
}
|
||
m := len(costMatrix[0])
|
||
|
||
used := make([]bool, m)
|
||
result := make([]int, n)
|
||
for i := range result {
|
||
result[i] = -1
|
||
}
|
||
|
||
for i := 0; i < n; i++ {
|
||
bestJ := -1
|
||
bestCost := 1.0
|
||
for j := 0; j < m; j++ {
|
||
if !used[j] && costMatrix[i][j] < bestCost {
|
||
bestCost = costMatrix[i][j]
|
||
bestJ = j
|
||
}
|
||
}
|
||
if bestJ >= 0 {
|
||
result[i] = bestJ
|
||
used[bestJ] = true
|
||
}
|
||
}
|
||
|
||
return result
|
||
}
|
||
|
||
// ===================== 工具函数 =====================
|
||
|
||
func decodeBase64ToImage(base64Str string) (image.Image, error) {
|
||
if strings.Contains(base64Str, ",") {
|
||
parts := strings.SplitN(base64Str, ",", 2)
|
||
if len(parts) == 2 {
|
||
base64Str = parts[1]
|
||
}
|
||
}
|
||
|
||
data, err := base64.StdEncoding.DecodeString(base64Str)
|
||
if err != nil {
|
||
data, err = base64.RawStdEncoding.DecodeString(base64Str)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
}
|
||
|
||
return imaging.Decode(bytes.NewReader(data))
|
||
} |