From 10a2acaf7b2df31ce7d7fc595028bd85954a8081 Mon Sep 17 00:00:00 2001 From: Admin Date: Fri, 17 Jul 2026 21:02:13 +0000 Subject: [PATCH] fix: use letterbox preprocessing for YOLO (same as ultralytics) - Keep aspect ratio with gray padding instead of stretching - Add ImageNet normalization (mean/std) - Proper coordinate restoration with offset/scale - This matches original Python ultralytics preprocessing --- internal/captcha/handler.go | 117 +++++++++++++++++++++++++++--------- 1 file changed, 89 insertions(+), 28 deletions(-) diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go index 14806f8..7fb3b89 100644 --- a/internal/captcha/handler.go +++ b/internal/captcha/handler.go @@ -873,8 +873,8 @@ func (h *Handler) detectYOLO(modelName, imageBase64 string) (*DetectionResult, e origW := img.Bounds().Dx() origH := img.Bounds().Dy() - // 预处理：Resize 到 640x640，归一化 - input, err := preprocessYOLO(img) + // 预处理：Letterbox resize 到 640x640，保持纵横比 + input, letterboxInfo, err := preprocessYOLO(img) if err != nil { return nil, err } @@ -899,7 +899,7 @@ func (h *Handler) detectYOLO(modelName, imageBase64 string) (*DetectionResult, e if origW < 300 || origH < 150 { confThresh = 0.05 // 小图使用更低阈值 } - detections := postprocessYOLO(output, origW, origH, confThresh, 0.45) + detections := postprocessYOLO(output, letterboxInfo, confThresh, 0.45) result := &DetectionResult{ Detections: make([]Detection, len(detections)), @@ -925,28 +925,84 @@ type YOLODetection struct { ClassName string } -// preprocessYOLO 预处理图片为 YOLO 输入 -func preprocessYOLO(img image.Image) ([]float32, error) { - // Resize 到 640x640 - resized := imaging.Resize(img, 640, 640, imaging.Lanczos) +// LetterboxInfo 保存 letterbox 预处理的参数，用于后处理坐标还原 +type LetterboxInfo struct { + OrigW, OrigH int // 原图尺寸 + Scale float64 // 缩放比例 + OffsetX int // X 方向偏移 + OffsetY int // Y 方向偏移 +} - 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 +// preprocessYOLO 预处理图片为 YOLO 输入 +// 使用 Letterbox 方式保持纵横比（与 ultralytics 一致） +func preprocessYOLO(img image.Image) ([]float32, LetterboxInfo, error) { + bounds := img.Bounds() + srcW, srcH := bounds.Dx(), bounds.Dy() + targetSize := 640 + + // Letterbox: 保持纵横比缩放，灰色填充 + scale := minFloat(float64(targetSize)/float64(srcW), float64(targetSize)/float64(srcH)) + newW := int(float64(srcW) * scale) + newH := int(float64(srcH) * scale) + + // 确保至少 1 像素 + if newW <= 0 { + newW = 1 + } + if newH <= 0 { + newH = 1 + } + + // 缩放图片 + resized := imaging.Resize(img, newW, newH, imaging.Lanczos) + + // 创建灰色背景画布 (128, 128, 128) + canvas := image.NewRGBA(image.Rect(0, 0, targetSize, targetSize)) + gray := color.RGBA{128, 128, 128, 255} + for x := 0; x < targetSize; x++ { + for y := 0; y < targetSize; y++ { + canvas.Set(x, y, gray) } } - return pixels, nil + // 居中粘贴缩放后的图片 + offsetX := (targetSize - newW) / 2 + offsetY := (targetSize - newH) / 2 + for x := 0; x < newW; x++ { + for y := 0; y < newH; y++ { + canvas.Set(offsetX+x, offsetY+y, resized.At(x, y)) + } + } + + // YOLO 标准化参数 (ImageNet) + mean := [3]float32{0.485, 0.456, 0.406} + std := [3]float32{0.229, 0.224, 0.225} + + pixels := make([]float32, 3*targetSize*targetSize) + for y := 0; y < targetSize; y++ { + for x := 0; x < targetSize; x++ { + c := canvas.At(x, y) + r, g, b, _ := c.RGBA() + // 归一化到 [0, 1]，然后标准化 + pixels[0*targetSize*targetSize+y*targetSize+x] = (float32(r)/65535.0 - mean[0]) / std[0] + pixels[1*targetSize*targetSize+y*targetSize+x] = (float32(g)/65535.0 - mean[1]) / std[1] + pixels[2*targetSize*targetSize+y*targetSize+x] = (float32(b)/65535.0 - mean[2]) / std[2] + } + } + + info := LetterboxInfo{ + OrigW: srcW, + OrigH: srcH, + Scale: scale, + OffsetX: offsetX, + OffsetY: offsetY, + } + + return pixels, info, nil } // postprocessYOLO YOLO11 后处理 -func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh float64) []YOLODetection { +func postprocessYOLO(output []float32, info LetterboxInfo, confThresh, nmsThresh float64) []YOLODetection { // YOLO11 ONNX 输出格式: [1, 5, 8400] 或 [1, 84, 8400] // 布局: [batch, features, boxes] // 对于单类检测 [1, 5, 8400]: @@ -972,10 +1028,6 @@ func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh f // 类别名称 classNames := []string{"icon"} - // 坐标缩放比例：从 640x640 到原图 - scaleX := float64(origW) / 640.0 - scaleY := float64(origH) / 640.0 - detections := []YOLODetection{} maxConf := 0.0 @@ -994,17 +1046,20 @@ func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh f continue } - // 获取边界框 (x, y, w, h) + // 获取边界框 (x, y, w, h) - 640x640 坐标系 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 + // Letterbox 坐标还原： + // 1. 减去 offset（灰色填充区域） + // 2. 除以 scale（还原到原图尺寸） + // 注意：offset 和 scale 对应的是 letterbox 参数 + cx = (cx - float64(info.OffsetX)) / info.Scale + cy = (cy - float64(info.OffsetY)) / info.Scale + w = w / info.Scale + h = h / info.Scale // 转换为 x1, y1, x2, y2 x1 := cx - w/2 @@ -1012,6 +1067,12 @@ func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh f x2 := cx + w/2 y2 := cy + h/2 + // 边界检查 + x1 = max(0, x1) + y1 = max(0, y1) + x2 = min(float64(info.OrigW), x2) + y2 = min(float64(info.OrigH), y2) + className := "object" if 0 < len(classNames) { className = classNames[0]