feat: implement YOLO detection using ONNX Runtime (OpenCV 4.6 incompatible with YOLO11)

internal/captcha/handler.go

@@ -98,22 +98,22 @@ func (h *Handler) loadModels() {
 		}
 	}
 
-	// 加载 YOLO 检测模型（YOLO11 需要特殊后处理，OpenCV 4.6 不完全支持）
+	// 加载 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 := opencv.LoadYOLO("icon", iconPath, ""); err != nil {
+		if err := onnx.LoadModel("detection_icon", iconPath); err != nil {
-			fmt.Printf("提示: Icon 检测模型暂不支持 (需要 YOLO11 后处理): %v\n", err)
+			fmt.Printf("提示: Icon 检测模型加载失败: %v\n", err)
 		} else {
-			fmt.Println("Icon 检测模型加载成功")
+			fmt.Println("Icon 检测模型加载成功 (ONNX Runtime)")
 		}
 	}
 
 	textPath := filepath.Join(h.modelPath, "Detection_Text.onnx")
 	if _, err := os.Stat(textPath); err == nil {
-		if err := opencv.LoadYOLO("text", textPath, ""); err != nil {
+		if err := onnx.LoadModel("detection_text", textPath); err != nil {
-			fmt.Printf("提示: Text 检测模型暂不支持 (需要 YOLO11 后处理): %v\n", err)
+			fmt.Printf("提示: Text 检测模型加载失败: %v\n", err)
 		} else {
-			fmt.Println("Text 检测模型加载成功")
+			fmt.Println("Text 检测模型加载成功 (ONNX Runtime)")
 		}
 	}
 }
@@ -858,17 +858,33 @@ func (h *Handler) DetectionText(imageBase64 string) (*DetectionResult, error) {
 }
 
 func (h *Handler) detectYOLO(modelName, imageBase64 string) (*DetectionResult, error) {
-	img, err := opencv.DecodeFromBase64(imageBase64)
+	// YOLO11 模型使用 ONNX Runtime（OpenCV 4.6 DNN 不兼容 YOLO11 输出格式）
-	if err != nil {
+	sess, ok := onnx.GetSession("detection_" + modelName)
-		return nil, err
+	if !ok {
+		return nil, fmt.Errorf("YOLO 检测模型 '%s' 未加载（需要 YOLO11 ONNX 模型）", modelName)
 	}
-	defer img.Free()
 
-	detections, err := opencv.DetectYOLO(modelName, img)
+	img, err := decodeBase64ToImage(imageBase64)
 	if err != nil {
 		return nil, err
 	}
 
+	// 预处理：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
+	}
+
+	// YOLO11 输出: [1, 5, 8400] 或 [1, 84, 8400]
+	// 需要转置 + NMS 后处理
+	detections := postprocessYOLO(output, 640, 640, 0.25, 0.45)
+
 	result := &DetectionResult{
 		Detections: make([]Detection, len(detections)),
 	}
@@ -883,6 +899,198 @@ func (h *Handler) detectYOLO(modelName, imageBase64 string) (*DetectionResult, e
 	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, imgW, imgH int, confThresh, nmsThresh float64) []YOLODetection {
+	// YOLO11 输出格式: [1, 5, 8400] 或 [1, 84, 8400]
+	// 5 = x, y, w, h, conf (单类)
+	// 84 = x, y, w, h, conf*80 (80类)
+	
+	// 检测输出维度
+	outputLen := len(output)
+	if outputLen == 0 {
+		return nil
+	}
+
+	// 假设输出是 [1, num_features, num_boxes]
+	// 尝试推断形状
+	numFeatures := 5 // 默认单类检测
+	numBoxes := 8400
+
+	if outputLen == numFeatures*numBoxes {
+		// [5, 8400] 格式
+	} else if outputLen%8400 == 0 {
+		numFeatures = outputLen / 8400
+	} else {
+		return nil
+	}
+
+	// 类别名称（根据模型调整）
+	classNames := []string{"icon"} // Detection_Icon 只有一个类
+
+	detections := []YOLODetection{}
+
+	// 遍历所有检测框
+	for i := 0; i < numBoxes; i++ {
+		// YOLO11 输出是 [features, boxes]，需要转置访问
+		// output[feat_idx * numBoxes + box_idx]
+		
+		// 获取置信度（第 5 个特征，索引 4）
+		confIdx := 4*numBoxes + i
+		if confIdx >= len(output) {
+			continue
+		}
+		confidence := float64(output[confIdx])
+
+		if confidence < confThresh {
+			continue
+		}
+
+		// 获取边界框 (x, y, w, h)
+		xIdx := 0*numBoxes + i
+		yIdx := 1*numBoxes + i
+		wIdx := 2*numBoxes + i
+		hIdx := 3*numBoxes + i
+
+		if xIdx >= len(output) || yIdx >= len(output) || wIdx >= len(output) || hIdx >= len(output) {
+			continue
+		}
+
+		cx := float64(output[xIdx])
+		cy := float64(output[yIdx])
+		w := float64(output[wIdx])
+		h := float64(output[hIdx])
+
+		// 转换为 x1, y1, x2, y2
+		x1 := cx - w/2
+		y1 := cy - h/2
+		x2 := cx + w/2
+		y2 := cy + h/2
+
+		classID := 0
+		className := "object"
+		if classID < len(classNames) {
+			className = classNames[classID]
+		}
+
+		detections = append(detections, YOLODetection{
+			X1:         x1,
+			Y1:         y1,
+			X2:         x2,
+			Y2:         y2,
+			Confidence: confidence,
+			ClassID:    classID,
+			ClassName:  className,
+		})
+	}
+
+	// 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) {