From 3ef9742376f211ea59e118205d2d7f40dddd9354 Mon Sep 17 00:00:00 2001
From: Admin <admin@viaeon.com>
Date: Fri, 17 Jul 2026 20:18:14 +0000
Subject: [PATCH] fix: YOLO output format is [boxes, features] not [features,
 boxes]

---
 internal/captcha/handler.go | 57 ++++++++++++++++++++++++-------------
 1 file changed, 37 insertions(+), 20 deletions(-)

diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go
index f8a417d..7ee8956 100644
--- a/internal/captcha/handler.go
+++ b/internal/captcha/handler.go
@@ -947,26 +947,38 @@ func preprocessYOLO(img image.Image) ([]float32, error) {
 
 // postprocessYOLO YOLO11 后处理
 func postprocessYOLO(output []float32, origW, origH 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类)
+	// YOLO11 输出格式可能是：
+	// [1, 5, 8400] - features 在前，boxes 在后
+	// [1, 8400, 5] - boxes 在前，features 在后（更常见）
 	
-	// 检测输出维度
 	outputLen := len(output)
 	if outputLen == 0 {
 		return nil
 	}
 
-	// 假设输出是 [1, num_features, num_boxes]
-	// 尝试推断形状
-	numFeatures := 5 // 默认单类检测
-	numBoxes := 8400
-
-	if outputLen == numFeatures*numBoxes {
-		// [5, 8400] 格式
+	// 尝试检测输出格式
+	// 如果是 [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)
 		return nil
 	}
 
@@ -983,12 +995,17 @@ func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh f
 	maxConf := 0.0
 
 	// 遍历所有检测框
+	// 格式 [boxes, features]: output[box_idx * numFeatures + feat_idx]
 	for i := 0; i < numBoxes; i++ {
-		// YOLO11 输出是 [features, boxes]，需要转置访问
-		// output[feat_idx * numBoxes + box_idx]
+		// 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
 		
-		// 获取置信度（第 5 个特征，索引 4）
-		confIdx := 4*numBoxes + i
+		// 获取置信度
+		confIdx := i * numFeatures + 4
 		if confIdx >= len(output) {
 			continue
 		}
@@ -1003,10 +1020,10 @@ func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh f
 		}
 
 		// 获取边界框 (x, y, w, h)
-		xIdx := 0*numBoxes + i
-		yIdx := 1*numBoxes + i
-		wIdx := 2*numBoxes + i
-		hIdx := 3*numBoxes + i
+		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
@@ -1046,7 +1063,7 @@ func postprocessYOLO(output []float32, origW, origH int, confThresh, nmsThresh f
 		})
 	}
 
-	fmt.Printf("DEBUG YOLO postprocess: outputLen=%d, numFeatures=%d, maxConf=%.4f, found=%d\n", outputLen, numFeatures, maxConf, len(detections))
+	fmt.Printf("DEBUG YOLO postprocess: outputLen=%d, format=[%d,%d], maxConf=%.4f, found=%d\n", outputLen, numBoxes, numFeatures, maxConf, len(detections))
 
 	// NMS
 	return nms(detections, nmsThresh)