feat: 完整实现 CGO 版本 - ONNX Runtime + OpenCV DNN (YOLO)

internal/captcha/handler.go

@@ -4,12 +4,15 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/base64"
+	"encoding/json"
 	"fmt"
 	"image"
+	"image/color"
 	"math"
 	"net/http"
 	"os"
 	"path/filepath"
+	"sort"
 	"strings"
 	"sync"
 
@@ -31,6 +34,8 @@ var modelConfigs = map[string]string{
 	"Rotation-RotNetR.onnx":  "[AntiCAP]-Rotation-RotNetR.onnx",
 	"Siamese-ResNet18.onnx":  "[AntiCAP]-Siamese-ResNet18.onnx",
 	"CharSets.txt":           "[Dddd]-CharSets.txt",
+	"Detection_Icon.onnx":    "Detection_Icon.onnx",
+	"Detection_Text.onnx":    "Detection_Text.onnx",
 }
 
 // 从 Gitea 仓库下载（公开仓库，无需认证）
@@ -88,11 +93,29 @@ func (h *Handler) loadModels() {
 			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
@@ -143,48 +166,43 @@ func (h *Handler) OCR(imageBase64 string) (string, error) {
 		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
 }
 
-// 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
+	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++ {
@@ -198,7 +216,6 @@ func preprocessOCR(img image.Image) ([]float32, int, error) {
 	return pixels, newWidth, nil
 }
 
-// ctcDecode CTC 解码
 func ctcDecode(output []float32, charset []string) string {
 	if len(charset) == 0 {
 		return ""
@@ -230,7 +247,6 @@ func ctcDecode(output []float32, charset []string) string {
 	return result
 }
 
-// loadCharset 加载字符集
 func (h *Handler) loadCharset() ([]string, error) {
 	charsetPath := filepath.Join(h.modelPath, "CharSets.txt")
 	file, err := os.Open(charsetPath)
@@ -248,7 +264,6 @@ func (h *Handler) loadCharset() ([]string, error) {
 		}
 	}
 
-	// 添加空白符作为第一个字符
 	result := make([]string, len(charset)+1)
 	result[0] = ""
 	copy(result[1:], charset)
@@ -266,32 +281,27 @@ func (h *Handler) Math(imageBase64 string) (string, error) {
 		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
@@ -300,21 +310,15 @@ func (h *Handler) Math(imageBase64 string) (string, error) {
 	return fmt.Sprintf("%v", result), nil
 }
 
-// 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
@@ -324,7 +328,6 @@ func preprocessMath(img image.Image) ([]float32, error) {
 	return pixels, nil
 }
 
-// decodeMath 解码数学表达式
 func decodeMath(output []float32) string {
 	numChars := len(mathChars) + 1
 	timesteps := len(output) / numChars
@@ -355,16 +358,12 @@ func decodeMath(output []float32) string {
 	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
@@ -391,7 +390,6 @@ func evalMathExpression(expr string) (interface{}, error) {
 		}
 	}
 
-	// 处理最后一个数字
 	switch op {
 	case '+':
 		result += num
@@ -405,7 +403,6 @@ func evalMathExpression(expr string) (interface{}, error) {
 		}
 	}
 
-	// 返回整数或浮点数
 	if result == float64(int(result)) {
 		return int(result), nil
 	}
@@ -414,48 +411,68 @@ func evalMathExpression(expr string) (interface{}, error) {
 
 // ===================== 滑块匹配 =====================
 
-func (h *Handler) SliderMatch(targetBase64, backgroundBase64 string) (int, error) {
-	target, err := opencv.DecodeFromBase64(targetBase64)
-	if err != nil {
-		return 0, err
-	}
-	defer target.Free()
-
-	background, err := opencv.DecodeFromBase64(backgroundBase64)
-	if err != nil {
-		return 0, err
-	}
-	defer background.Free()
-
-	return opencv.SliderMatch(target, background)
+type SliderMatchResult struct {
+	Target []int `json:"target"`
 }
 
-func (h *Handler) SliderComparison(targetBase64, backgroundBase64 string) (int, error) {
+func (h *Handler) SliderMatch(targetBase64, backgroundBase64 string) (*SliderMatchResult, error) {
 	target, err := opencv.DecodeFromBase64(targetBase64)
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
 	defer target.Free()
 
 	background, err := opencv.DecodeFromBase64(backgroundBase64)
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
 	defer background.Free()
 
-	return opencv.SliderComparison(target, background)
+	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) {
-	// 使用 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
@@ -482,7 +499,6 @@ func (h *Handler) compareSimilarityONNX(sess *onnx.Session, img1Base64, img2Base
 		return 0, err
 	}
 
-	// 预处理
 	input1, err := preprocessSiamese(img1)
 	if err != nil {
 		return 0, err
@@ -493,19 +509,16 @@ func (h *Handler) compareSimilarityONNX(sess *onnx.Session, img1Base64, img2Base
 		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]
@@ -513,7 +526,6 @@ func (h *Handler) compareSimilarityONNX(sess *onnx.Session, img1Base64, img2Base
 		}
 		dist = float32(math.Sqrt(float64(dist)))
 
-		// 相似度
 		similarity := 1.0 / (1.0 + dist)
 		return similarity, nil
 	}
@@ -522,11 +534,9 @@ func (h *Handler) compareSimilarityONNX(sess *onnx.Session, img1Base64, img2Base
 }
 
 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}
 
@@ -546,43 +556,43 @@ func preprocessSiamese(img image.Image) ([]float32, error) {
 
 // ===================== 旋转检测 =====================
 
-func (h *Handler) SingleRotate(imageBase64 string) (float32, error) {
-	// 使用 ONNX 模型
+func (h *Handler) SingleRotate(imageBase64 string) (int, error) {
 	sess, ok := onnx.GetSession("rotate")
 	if ok {
 		return h.singleRotateONNX(sess, imageBase64)
 	}
 
-	// 使用 OpenCV
 	img, err := opencv.DecodeFromBase64(imageBase64)
 	if err != nil {
 		return 0, err
 	}
 	defer img.Free()
 
-	return opencv.DetectRotation(img)
+	angle, err := opencv.DetectRotation(img)
+	if err != nil {
+		return 0, err
+	}
+
+	return int(angle), nil
 }
 
-func (h *Handler) singleRotateONNX(sess *onnx.Session, imageBase64 string) (float32, error) {
+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++ {
@@ -592,15 +602,30 @@ func (h *Handler) singleRotateONNX(sess *onnx.Session, imageBase64 string) (floa
 		}
 	}
 
-	return float32(maxIdx), nil
+	return maxIdx, nil
 }
 
 func preprocessRotation(img image.Image) ([]float32, error) {
-	// 调整大小为 224x224
-	resized := imaging.Resize(img, 224, 224, imaging.Lanczos)
+	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)
 
-	// ImageNet 归一化
 	mean := [3]float32{0.485, 0.456, 0.406}
 	std := [3]float32{0.229, 0.224, 0.225}
 
@@ -618,66 +643,311 @@ func preprocessRotation(img image.Image) ([]float32, error) {
 	return pixels, nil
 }
 
-func (h *Handler) DoubleRotate(insideBase64, outsideBase64 string) (float32, error) {
+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 0, err
+		return nil, err
 	}
 	defer inside.Free()
 
 	outside, err := opencv.DecodeFromBase64(outsideBase64)
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
 	defer outside.Free()
 
 	angleInside, err := opencv.DetectRotation(inside)
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
 
 	angleOutside, err := opencv.DetectRotation(outside)
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
 
-	return angleInside - angleOutside, nil
+	angle := int(angleInside - angleOutside)
+	if anticlockwise {
+		angle = -angle
+	}
+	if angle < 0 {
+		angle += 360
+	}
+
+	return &DoubleRotateResult{Angle: angle}, nil
 }
 
-// ===================== 图标/文字检测 =====================
+// ===================== 图标/文字检测 (YOLO via OpenCV DNN) =====================
 
-func (h *Handler) DetectionIcon(imageBase64 string) ([]map[string]int, error) {
-	// 暂时返回空结果
-	return []map[string]int{}, nil
+type Detection struct {
+	Class string `json:"class"`
+	Box   []int  `json:"box"`
 }
 
-func (h *Handler) DetectionText(imageBase64 string) ([]map[string]int, error) {
-	return []map[string]int{}, nil
+type DetectionResult struct {
+	Detections []Detection `json:"detections"`
 }
 
-func (h *Handler) DetectionIconOrder(orderImgBase64, targetImgBase64 string) ([]map[string]int, error) {
-	return []map[string]int{}, nil
+func (h *Handler) DetectionIcon(imageBase64 string) (*DetectionResult, error) {
+	return h.detectYOLO("icon", imageBase64)
 }
 
-func (h *Handler) DetectionTextOrder(orderImgBase64, targetImgBase64 string) ([]map[string]int, error) {
-	return []map[string]int{}, nil
+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 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 {
+		data, err = base64.RawStdEncoding.DecodeString(base64Str)
+		if err != nil {
 			return nil, err
 		}
 	}