diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go index ca02b93..0d445a1 100644 --- a/internal/captcha/handler.go +++ b/internal/captcha/handler.go @@ -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 } } diff --git a/test_full_rotation.go b/test_full_rotation.go new file mode 100644 index 0000000..8a2868b --- /dev/null +++ b/test_full_rotation.go @@ -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 +}