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
This commit is contained in:
+87
-26
@@ -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
|
||||
}
|
||||
|
||||
// LetterboxInfo 保存 letterbox 预处理的参数,用于后处理坐标还原
|
||||
type LetterboxInfo struct {
|
||||
OrigW, OrigH int // 原图尺寸
|
||||
Scale float64 // 缩放比例
|
||||
OffsetX int // X 方向偏移
|
||||
OffsetY int // Y 方向偏移
|
||||
}
|
||||
|
||||
// preprocessYOLO 预处理图片为 YOLO 输入
|
||||
func preprocessYOLO(img image.Image) ([]float32, error) {
|
||||
// Resize 到 640x640
|
||||
resized := imaging.Resize(img, 640, 640, imaging.Lanczos)
|
||||
// 使用 Letterbox 方式保持纵横比(与 ultralytics 一致)
|
||||
func preprocessYOLO(img image.Image) ([]float32, LetterboxInfo, error) {
|
||||
bounds := img.Bounds()
|
||||
srcW, srcH := bounds.Dx(), bounds.Dy()
|
||||
targetSize := 640
|
||||
|
||||
pixels := make([]float32, 3*640*640)
|
||||
for y := 0; y < 640; y++ {
|
||||
for x := 0; x < 640; x++ {
|
||||
c := resized.At(x, y)
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
|
||||
// 居中粘贴缩放后的图片
|
||||
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*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
|
||||
// 归一化到 [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]
|
||||
}
|
||||
}
|
||||
|
||||
return pixels, nil
|
||||
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]
|
||||
|
||||
Reference in New Issue
Block a user