diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go index 7ee8956..0956d37 100644 --- a/internal/captcha/handler.go +++ b/internal/captcha/handler.go @@ -947,68 +947,43 @@ func preprocessYOLO(img image.Image) ([]float32, error) { // postprocessYOLO YOLO11 后处理 func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh float64) []YOLODetection { - // YOLO11 输出格式可能是: - // [1, 5, 8400] - features 在前,boxes 在后 - // [1, 8400, 5] - boxes 在前,features 在后(更常见) + // YOLO11 ONNX 输出格式: [1, 5, 8400] 或 [1, 84, 8400] + // 布局: [batch, features, boxes] + // 对于单类检测 [1, 5, 8400]: + // output[0:8400] = 所有 box 的 x (中心点) + // output[8400:16800] = 所有 box 的 y (中心点) + // output[16800:25200] = 所有 box 的 w + // output[25200:33600] = 所有 box 的 h + // output[33600:42000] = 所有 box 的 confidence outputLen := len(output) if outputLen == 0 { return nil } - // 尝试检测输出格式 - // 如果是 [5, 8400],则 output[4*8400+0] 是第一个 box 的置信度 - // 如果是 [8400, 5],则 output[0*5+4] 是第一个 box 的置信度 - - // 检查第一个值是否像坐标(>10)还是概率(<10) - // 从输出值看,第一个值 ~17,说明是坐标而非概率 - // 所以格式应该是 [8400, 5] - - var numBoxes, numFeatures int - if outputLen == 5*8400 { - // 尝试两种格式 - // 格式1: [5, 8400] - 第 5 个元素是 output[4] - // 格式2: [8400, 5] - 第 5 个元素是 output[4] - - // 从日志看,output[0]~output[19] 都是 17-154 范围,像是坐标 - // 这说明格式是 [boxes, features],因为 output[0] 是第一个 box 的 x - numBoxes = 8400 - numFeatures = 5 - } else if outputLen%8400 == 0 { - numFeatures = outputLen / 8400 - numBoxes = 8400 - } else { - fmt.Printf("DEBUG YOLO: unknown output format, len=%d\n", outputLen) + numBoxes := 8400 + numFeatures := 5 + + if outputLen != numFeatures*numBoxes { + fmt.Printf("DEBUG YOLO: unexpected output length %d (expected %d)\n", outputLen, numFeatures*numBoxes) return nil } - // 类别名称(根据模型调整) - classNames := []string{"icon"} // Detection_Icon 只有一个类 + // 类别名称 + classNames := []string{"icon"} // 坐标缩放比例:从 640x640 到原图 scaleX := float64(origW) / 640.0 scaleY := float64(origH) / 640.0 detections := []YOLODetection{} - - // 记录最大置信度用于调试 maxConf := 0.0 // 遍历所有检测框 - // 格式 [boxes, features]: output[box_idx * numFeatures + feat_idx] for i := 0; i < numBoxes; i++ { - // YOLO11 输出: [boxes, features] - // output[i*5 + 0] = x - // output[i*5 + 1] = y - // output[i*5 + 2] = w - // output[i*5 + 3] = h - // output[i*5 + 4] = confidence - - // 获取置信度 - confIdx := i * numFeatures + 4 - if confIdx >= len(output) { - continue - } + // [features, boxes] 布局 + // output[feat_idx * numBoxes + box_idx] + confIdx := 4*numBoxes + i // 第 5 个特征(索引 4) confidence := float64(output[confIdx]) if confidence > maxConf { @@ -1020,19 +995,10 @@ func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh f } // 获取边界框 (x, y, w, h) - xIdx := i * numFeatures + 0 - yIdx := i * numFeatures + 1 - wIdx := i * numFeatures + 2 - hIdx := i * numFeatures + 3 - - 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]) + 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 @@ -1046,10 +1012,9 @@ func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh f x2 := cx + w/2 y2 := cy + h/2 - classID := 0 className := "object" - if classID < len(classNames) { - className = classNames[classID] + if 0 < len(classNames) { + className = classNames[0] } detections = append(detections, YOLODetection{ @@ -1058,12 +1023,12 @@ func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh f X2: x2, Y2: y2, Confidence: confidence, - ClassID: classID, + ClassID: 0, ClassName: className, }) } - fmt.Printf("DEBUG YOLO postprocess: outputLen=%d, format=[%d,%d], maxConf=%.4f, found=%d\n", outputLen, numBoxes, numFeatures, maxConf, len(detections)) + fmt.Printf("DEBUG YOLO postprocess: outputLen=%d, format=[1,%d,%d], maxConf=%.4f, found=%d\n", outputLen, numFeatures, numBoxes, maxConf, len(detections)) // NMS return nms(detections, nmsThresh)