package main
import (
"fmt"
"image"
"io/ioutil"
"os"
tf "github.com/tensorflow/tensorflow/tensorflow/go"
"github.com/nfnt/resize"
_ "image/png"
)
func preprocessImage(imagePath string) (*tf.Tensor, error) {
file, err := os.Open(imagePath)
if err != nil {
return nil, err
}
defer file.Close()
img, _, err := image.Decode(file)
if err != nil {
return nil, err
}
resized := resize.Resize(160, 60, img, resize.Lanczos3)
// 转换为灰度 + float32 + 正则化
data := make([]float32, 160*60)
i := 0
for y := 0; y < 60; y++ {
for x := 0; x < 160; x++ {
r, g, b, _ := resized.At(x, y).RGBA()
gray := float32((r + g + b) / 3 >> 8)
data[i] = gray / 255.0
i++
}
}
// 构建 Tensor
tensor, err := tf.NewTensor([1][60][160][1]float32{})
if err != nil {
return nil, err
}
copy(tensor.Value().([1][60][160][1]float32)[0][:][0][:], data)
return tensor, nil
}
func main() {
model, err := tf.LoadSavedModel("saved_model/captcha_model", []string{"serve"}, nil)
if err != nil {
panic(err)
}
defer model.Session.Close()
input, err := preprocessImage("captcha_samples/A7K9_1.png")
if err != nil {
panic(err)
}
result, err := model.Session.Run(
map[tf.Output]*tf.Tensor{
model.Graph.Operation("serving_default_input_1").Output(0): input,
},
[]tf.Output{
model.Graph.Operation("StatefulPartitionedCall").Output(0),
},
nil,
)
if err != nil {
panic(err)
}
// 输出预测
prediction := result[0].Value().([][]float32)[0]
alphabet := "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
fmt.Print("Predicted: ")
for i := 0; i < 4; i++ {
// 每个字符有36个概率值
offset := i * 36
maxIdx := 0
maxVal := float32(0)
for j := 0; j < 36; j++ {
if prediction[offset+j] > maxVal {
maxVal = prediction[offset+j]
maxIdx = j
}
}
fmt.Print(string(alphabet[maxIdx]))
}
fmt.Println()
}
- 1.
- 2.
- 3.
- 4.
- 5.
- 6.
- 7.
- 8.
- 9.
- 10.
- 11.
- 12.
- 13.
- 14.
- 15.
- 16.
- 17.
- 18.
- 19.
- 20.
- 21.
- 22.
- 23.
- 24.
- 25.
- 26.
- 27.
- 28.
- 29.
- 30.
- 31.
- 32.
- 33.
- 34.
- 35.
- 36.
- 37.
- 38.
- 39.
- 40.
- 41.
- 42.
- 43.
- 44.
- 45.
- 46.
- 47.
- 48.
- 49.
- 50.
- 51.
- 52.
- 53.
- 54.
- 55.
- 56.
- 57.
- 58.
- 59.
- 60.
- 61.
- 62.
- 63.
- 64.
- 65.
- 66.
- 67.
- 68.
- 69.
- 70.
- 71.
- 72.
- 73.
- 74.
- 75.
- 76.
- 77.
- 78.
- 79.
- 80.
- 81.
- 82.
- 83.
- 84.
- 85.
- 86.
- 87.
- 88.
- 89.
- 90.
- 91.