960 lines
22 KiB
Go
960 lines
22 KiB
Go
package captcha
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"image"
|
||
"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_Icon.onnx": {remote: "Detection_Icon.onnx", optional: false},
|
||
"Detection_Text.onnx": {remote: "Detection_Text.onnx", optional: false},
|
||
}
|
||
|
||
// 从 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 检测模型
|
||
iconPath := filepath.Join(h.modelPath, "Detection_Icon.onnx")
|
||
if _, err := os.Stat(iconPath); err == nil {
|
||
if err := opencv.LoadYOLO("icon", iconPath, ""); err != nil {
|
||
fmt.Printf("警告: 加载 Icon 检测模型失败: %v\n", err)
|
||
} else {
|
||
fmt.Println("Icon 检测模型加载成功")
|
||
}
|
||
}
|
||
|
||
textPath := filepath.Join(h.modelPath, "Detection_Text.onnx")
|
||
if _, err := os.Stat(textPath); err == nil {
|
||
if err := opencv.LoadYOLO("text", textPath, ""); err != nil {
|
||
fmt.Printf("警告: 加载 Text 检测模型失败: %v\n", err)
|
||
} else {
|
||
fmt.Println("Text 检测模型加载成功")
|
||
}
|
||
}
|
||
}
|
||
|
||
// 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) (string, error) {
|
||
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
|
||
}
|
||
|
||
func preprocessMath(img image.Image) ([]float32, error) {
|
||
resized := imaging.Resize(img, 200, 70, imaging.Lanczos)
|
||
rgb := imaging.Clone(resized)
|
||
|
||
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()
|
||
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 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
|
||
}
|
||
|
||
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) {
|
||
img, err := opencv.DecodeFromBase64(imageBase64)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer img.Free()
|
||
|
||
detections, err := opencv.DetectYOLO(modelName, img)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
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)},
|
||
}
|
||
}
|
||
|
||
return result, nil
|
||
}
|
||
|
||
// ===================== 按序点击 (匈牙利算法匹配) =====================
|
||
|
||
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))
|
||
} |