fix: remove ImageNet normalization for YOLO (match Ultralytics default)
- Ultralytics YOLO only does /255 normalization, no ImageNet mean/std - This was causing detection box coordinate drift
This commit is contained in:
@@ -982,19 +982,17 @@ func preprocessYOLO(img image.Image) ([]float32, LetterboxInfo, error) {
|
||||
}
|
||||
}
|
||||
|
||||
// YOLO 标准化参数 (ImageNet)
|
||||
mean := [3]float32{0.485, 0.456, 0.406}
|
||||
std := [3]float32{0.229, 0.224, 0.225}
|
||||
|
||||
// YOLO 归一化:只需要 /255,不需要 ImageNet 标准化
|
||||
// Ultralytics YOLO 默认预处理:letterbox + /255
|
||||
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]
|
||||
// 归一化到 [0, 1]
|
||||
pixels[0*targetSize*targetSize+y*targetSize+x] = float32(r) / 65535.0
|
||||
pixels[1*targetSize*targetSize+y*targetSize+x] = float32(g) / 65535.0
|
||||
pixels[2*targetSize*targetSize+y*targetSize+x] = float32(b) / 65535.0
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
//go:build ignore
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"math"
|
||||
|
||||
"github.com/disintegration/imaging"
|
||||
"anticaptcha/pkg/onnx"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// 加载模型
|
||||
err := onnx.LoadModel("rotate", "models/[AntiCAP]-Rotation-RotNetR.onnx")
|
||||
if err != nil {
|
||||
fmt.Printf("Failed to load model: %v\n", err)
|
||||
return
|
||||
}
|
||||
fmt.Println("Model loaded")
|
||||
|
||||
// 创建测试图片 (224x224)
|
||||
img := image.NewRGBA(image.Rect(0, 0, 224, 224))
|
||||
for y := 0; y < 224; y++ {
|
||||
for x := 0; x < 224; x++ {
|
||||
img.Set(x, y, color.RGBA{100, 150, 200, 255})
|
||||
}
|
||||
}
|
||||
for y := 50; y < 150; y++ {
|
||||
for x := 50; x < 100; x++ {
|
||||
img.Set(x, y, color.RGBA{255, 0, 0, 255})
|
||||
}
|
||||
}
|
||||
for y := 80; y < 140; y++ {
|
||||
for x := 120; x < 180; x++ {
|
||||
cx, cy := 150.0, 110.0
|
||||
rx, ry := 30.0, 30.0
|
||||
dx, dy := float64(x)-cx, float64(y)-cy
|
||||
if (dx*dx)/(rx*rx)+(dy*dy)/(ry*ry) <= 1 {
|
||||
img.Set(x, y, color.RGBA{0, 255, 0, 255})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 预处理
|
||||
input := preprocessRotation(img)
|
||||
fmt.Printf("Input length: %d\n", len(input))
|
||||
|
||||
// 推理
|
||||
sess, ok := onnx.GetSession("rotate")
|
||||
if !ok {
|
||||
fmt.Println("Session not found")
|
||||
return
|
||||
}
|
||||
|
||||
dims := []int64{1, 3, 224, 224}
|
||||
output, err := sess.Run(input, dims)
|
||||
if err != nil {
|
||||
fmt.Printf("Run failed: %v\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Printf("Output length: %d\n", len(output))
|
||||
fmt.Printf("Output first 10: %v\n", output[:10])
|
||||
|
||||
// Argmax
|
||||
maxIdx := 0
|
||||
maxProb := float32(-math.MaxFloat32)
|
||||
for i := 0; i < len(output); i++ {
|
||||
if output[i] > maxProb {
|
||||
maxProb = output[i]
|
||||
maxIdx = i
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("\nPredicted angle: %d degrees\n", maxIdx)
|
||||
}
|
||||
|
||||
func preprocessRotation(img image.Image) []float32 {
|
||||
bounds := img.Bounds()
|
||||
w, h := bounds.Dx(), bounds.Dy()
|
||||
|
||||
size := w
|
||||
if h < w {
|
||||
size = h
|
||||
}
|
||||
|
||||
cropX := (w - size) / 2
|
||||
cropY := (h - size) / 2
|
||||
cropped := imaging.Crop(img, image.Rect(cropX, cropY, cropX+size, cropY+size))
|
||||
|
||||
sqrt2 := math.Sqrt(2.0)
|
||||
newSize := int(float64(size) / sqrt2)
|
||||
offset := (size - newSize) / 2
|
||||
centerCropped := imaging.Crop(cropped, image.Rect(offset, offset, offset+newSize, offset+newSize))
|
||||
|
||||
resized := imaging.Resize(centerCropped, 224, 224, imaging.Lanczos)
|
||||
rgb := imaging.Clone(resized)
|
||||
|
||||
mean := [3]float32{0.485, 0.456, 0.406}
|
||||
std := [3]float32{0.229, 0.224, 0.225}
|
||||
|
||||
pixels := make([]float32, 3*224*224)
|
||||
for y := 0; y < 224; y++ {
|
||||
for x := 0; x < 224; x++ {
|
||||
c := rgb.At(x, y)
|
||||
r, g, b, _ := c.RGBA()
|
||||
pixels[0*224*224+y*224+x] = (float32(r)/65535.0 - mean[0]) / std[0]
|
||||
pixels[1*224*224+y*224+x] = (float32(g)/65535.0 - mean[1]) / std[1]
|
||||
pixels[2*224*224+y*224+x] = (float32(b)/65535.0 - mean[2]) / std[2]
|
||||
}
|
||||
}
|
||||
|
||||
return pixels
|
||||
}
|
||||
Reference in New Issue
Block a user