admin
47970834ec
- Ultralytics YOLO only does /255 normalization, no ImageNet mean/std - This was causing detection box coordinate drift
119 lines
2.6 KiB
Go
119 lines
2.6 KiB
Go
//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
|
|
}
|