fix: YOLO output format is [boxes, features] not [features, boxes]
This commit is contained in:
+36
-19
@@ -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
|
||||
// 尝试检测输出格式
|
||||
// 如果是 [5, 8400],则 output[4*8400+0] 是第一个 box 的置信度
|
||||
// 如果是 [8400, 5],则 output[0*5+4] 是第一个 box 的置信度
|
||||
|
||||
if outputLen == numFeatures*numBoxes {
|
||||
// [5, 8400] 格式
|
||||
// 检查第一个值是否像坐标(>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)
|
||||
|
||||
Reference in New Issue
Block a user