使用 Go 构建图像验证码识别 REST API 服务(结合 TensorFlow Lite)
本项目介绍如何使用 Go 实现一个图像验证码识别的后端接口,使用 Python 训练好的模型(TFLite 格式)进行识别,并通过 HTTP 接口对外提供服务。
一、项目结构
captcha-api/
├── model/
│ └── captcha.tflite
├── main.go
├── handler.go
├── imageutil.go
├── go.mod
二、安装依赖
使用如下命令安装依赖:
更多内容访问ttocr.com或联系1436423940
go mod init captcha-api
go get github.com/gin-gonic/gin
go get github.com/mattn/go-tflite
go get github.com/nfnt/resize
三、图像处理工具(imageutil.go)
package main
import (
"image"
"image/jpeg"
"image/png"
"os"
"strings"
"github.com/nfnt/resize"
)
func LoadImageAsFloat32(filePath string, w, h int) ([]float32, error) {
f, err := os.Open(filePath)
if err != nil {
return nil, err
}
defer f.Close()
var img image.Image
if strings.HasSuffix(filePath, ".png") {
img, err = png.Decode(f)
} else {
img, err = jpeg.Decode(f)
}
if err != nil {
return nil, err
}
resized := resize.Resize(uint(w), uint(h), img, resize.Bilinear)
data := make([]float32, w*h*3)
idx := 0
for y := 0; y < h; y++ {
for x := 0; x < w; x++ {
r, g, b, _ := resized.At(x, y).RGBA()
data[idx] = float32(r>>8) / 255.0
data[idx+1] = float32(g>>8) / 255.0
data[idx+2] = float32(b>>8) / 255.0
idx += 3
}
}
return data, nil
}
四、识别处理函数(handler.go)
package main
import (
"fmt"
"net/http"
"os"
"path/filepath"
"github.com/gin-gonic/gin"
"github.com/mattn/go-tflite"
)
var charset = []rune("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ")
var model *tflite.Model
var interpreter *tflite.Interpreter
func InitModel() {
model = tflite.NewModelFromFile("model/captcha.tflite")
if model == nil {
panic("无法加载模型")
}
interpreter = tflite.NewInterpreter(model, tflite.NewInterpreterOptions())
if interpreter.AllocateTensors() != nil {
panic("张量初始化失败")
}
}
func Argmax(arr []float32) int {
maxIdx := 0
for i := 1; i < len(arr); i++ {
if arr[i] > arr[maxIdx] {
maxIdx = i
}
}
return maxIdx
}
func PredictImage(path string) (string, error) {
input, err := LoadImageAsFloat32(path, 160, 60)
if err != nil {
return "", err
}
interpreter.GetInputTensor(0).CopyFromBuffer(input)
if interpreter.Invoke() != nil {
return "", fmt.Errorf("推理失败")
}
result := ""
for i := 0; i < 4; i++ {
output := interpreter.GetOutputTensor(i)
buf := make([]float32, len(charset))
output.CopyToBuffer(&buf[0])
result += string(charset[Argmax(buf)])
}
return result, nil
}
func UploadHandler(c *gin.Context) {
file, err := c.FormFile("image")
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "未找到图像"})
return
}
tempPath := filepath.Join(os.TempDir(), file.Filename)
if err := c.SaveUploadedFile(file, tempPath); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存图像失败"})
return
}
pred, err := PredictImage(tempPath)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "识别失败"})
return
}
os.Remove(tempPath)
c.JSON(http.StatusOK, gin.H{"captcha": pred})
}
五、主函数入口(main.go)
package main
import (
"github.com/gin-gonic/gin"
)
func main() {
InitModel()
r := gin.Default()
r.POST("/recognize", UploadHandler)
r.Run(":8080")
}
六、使用示例(cURL)
上传验证码图像进行识别:
curl -X POST -F "image=@captcha_samples/AB12_001.png" http://localhost:8080/recognize
响应结果:
{"captcha":"AB12"}
浙公网安备 33010602011771号