使用 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"}

posted @ 2025-06-01 11:59  ttocr、com  阅读(17)  评论(0)    收藏  举报