From b0555e228131fd85cfc415745773e8a5b74af448 Mon Sep 17 00:00:00 2001 From: Admin Date: Fri, 17 Jul 2026 12:59:36 +0000 Subject: [PATCH] feat: implement YOLO detection using ONNX Runtime (OpenCV 4.6 incompatible with YOLO11) --- internal/captcha/handler.go | 232 ++++++++++++++++++++++++++++++++++-- 1 file changed, 220 insertions(+), 12 deletions(-) diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go index 46bf0e1..6e45917 100644 --- a/internal/captcha/handler.go +++ b/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 { - fmt.Printf("提示: Icon 检测模型暂不支持 (需要 YOLO11 后处理): %v\n", err) + if err := onnx.LoadModel("detection_icon", iconPath); err != nil { + 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 { - fmt.Printf("提示: Text 检测模型暂不支持 (需要 YOLO11 后处理): %v\n", err) + if err := onnx.LoadModel("detection_text", textPath); err != nil { + 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) - if err != nil { - return nil, err + // 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) } - 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) {