287 lines
7.0 KiB
Go
287 lines
7.0 KiB
Go
package backend
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"math"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/disintegration/imaging"
|
|
ort "github.com/yalue/onnxruntime_go"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var (
|
|
publicImagePath = "./frontend/public/images"
|
|
modelPath = "resnet_epoch_100.onnx"
|
|
|
|
// ImageNet 标准化参数
|
|
mean = []float32{0.485, 0.456, 0.406}
|
|
std = []float32{0.229, 0.224, 0.225}
|
|
|
|
// 猫品种标签
|
|
labelName = []string{
|
|
"american_shorthair",
|
|
"bengal",
|
|
"british_shorthair",
|
|
"exotic_shorthair",
|
|
"maine_coon",
|
|
"ragdoll",
|
|
"sphynx",
|
|
}
|
|
)
|
|
|
|
func NewApp() *App {
|
|
return &App{}
|
|
}
|
|
|
|
func (a *App) Startup(ctx context.Context) {
|
|
a.ctx = ctx
|
|
|
|
// 设置 ONNX Runtime DLL 路径
|
|
exeDir, _ := os.Executable()
|
|
println("exeDir", exeDir)
|
|
ort.SetSharedLibraryPath(filepath.Join(filepath.Dir(exeDir), "onnxruntime.dll"))
|
|
|
|
// 初始化 ONNX Runtime 环境
|
|
err := ort.InitializeEnvironment()
|
|
if err != nil {
|
|
panic("Failed to initialize ONNX runtime: " + err.Error())
|
|
}
|
|
}
|
|
|
|
func preprocessImage(imgData []byte) ([]float32, error) {
|
|
// 解码图片
|
|
reader := bytes.NewReader(imgData)
|
|
img, err := imaging.Decode(reader, imaging.AutoOrientation(true))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("decode image: %w", err)
|
|
}
|
|
|
|
// 缩放到 224x224
|
|
img = imaging.Resize(img, 224, 224, imaging.Lanczos)
|
|
|
|
// 转换为 float32 数组 (NCHW 格式: 1, 3, 224, 224)
|
|
input := make([]float32, 1*3*224*224)
|
|
|
|
bounds := img.Bounds()
|
|
idx := 0
|
|
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
|
|
for x := bounds.Min.X; x < bounds.Max.X; x++ {
|
|
r, g, b, _ := img.At(x, y).RGBA()
|
|
// RGBA 返回 0-65535，需要转换到 0-255
|
|
rf := float32(r>>8) / 255.0
|
|
gf := float32(g>>8) / 255.0
|
|
bf := float32(b>>8) / 255.0
|
|
|
|
// ImageNet 标准化
|
|
input[idx] = (rf - mean[0]) / std[0] // R channel
|
|
input[idx+224*224] = (gf - mean[1]) / std[1] // G channel
|
|
input[idx+224*224*2] = (bf - mean[2]) / std[2] // B channel
|
|
idx++
|
|
}
|
|
}
|
|
|
|
return input, nil
|
|
}
|
|
|
|
func runInference(input []float32) (string, float64, error) {
|
|
// 创建输入张量 [1, 3, 224, 224]
|
|
inputTensor, err := ort.NewTensor(ort.Shape{1, 3, 224, 224}, input)
|
|
if err != nil {
|
|
return "", 0, fmt.Errorf("create input tensor: %w", err)
|
|
}
|
|
defer inputTensor.Destroy()
|
|
|
|
// 创建输出张量 [1, 7]
|
|
outputTensor, err := ort.NewTensor(ort.Shape{1, 7}, make([]float32, 7))
|
|
if err != nil {
|
|
return "", 0, fmt.Errorf("create output tensor: %w", err)
|
|
}
|
|
defer outputTensor.Destroy()
|
|
|
|
exeDir, _ := os.Executable()
|
|
println("exeDir222", exeDir)
|
|
|
|
// 创建 session 并运行推理
|
|
session, err := ort.NewAdvancedSession(
|
|
filepath.Join(filepath.Dir(exeDir), modelPath),
|
|
[]string{"input"},
|
|
[]string{"output"},
|
|
[]ort.Value{inputTensor},
|
|
[]ort.Value{outputTensor},
|
|
nil,
|
|
)
|
|
if err != nil {
|
|
return "", 0, fmt.Errorf("create session: %w", err)
|
|
}
|
|
defer session.Destroy()
|
|
|
|
err = session.Run()
|
|
if err != nil {
|
|
return "", 0, fmt.Errorf("run inference: %w", err)
|
|
}
|
|
|
|
// 获取输出
|
|
outputData := outputTensor.GetData()
|
|
|
|
// 找最大值的索引
|
|
maxIdx := 0
|
|
maxVal := outputData[0]
|
|
for i := 1; i < len(outputData); i++ {
|
|
if outputData[i] > maxVal {
|
|
maxVal = outputData[i]
|
|
maxIdx = i
|
|
}
|
|
}
|
|
|
|
// Softmax 计算置信度
|
|
var sum float64
|
|
for _, v := range outputData {
|
|
sum += math.Exp(float64(v))
|
|
}
|
|
confidence := math.Exp(float64(maxVal)) / sum
|
|
|
|
return labelName[maxIdx], confidence, nil
|
|
}
|
|
|
|
func (a *App) GormDB() (*gorm.DB, error) {
|
|
db, err := gorm.Open(sqlite.Open("app.db"), &gorm.Config{})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return db, nil
|
|
}
|
|
|
|
func (a *App) UploadImage(data []byte, filename string) Response {
|
|
uploadsDir := publicImagePath
|
|
err := os.MkdirAll(uploadsDir, 0755)
|
|
if err != nil {
|
|
return Response{Code: 1, Message: "failed", Data: err.Error()}
|
|
}
|
|
|
|
ext := filepath.Ext(filename)
|
|
newFilename := strconv.FormatInt(time.Now().UnixMilli(), 10) + ext
|
|
filePath := filepath.Join(uploadsDir, newFilename)
|
|
|
|
err = os.WriteFile(filePath, data, 0644)
|
|
if err != nil {
|
|
return Response{Code: 1, Message: "failed", Data: err.Error()}
|
|
}
|
|
|
|
println("333")
|
|
|
|
return Response{Code: 0, Message: "success", Data: newFilename}
|
|
}
|
|
|
|
func (a *App) GetHistory(page int, pageSize int) Response {
|
|
db, err := a.GormDB()
|
|
if err != nil {
|
|
return Response{Code: 1, Message: err.Error()}
|
|
}
|
|
|
|
var total int64
|
|
db.Table("history_test").Count(&total)
|
|
|
|
var historyList []HistoryWithBreed
|
|
err = db.Table("history_test").Select("history_test.*, breeds_test.brief, breeds_test.name").
|
|
Joins("LEFT JOIN breeds_test ON history_test.breed = breeds_test.id").
|
|
Order("history_test.id DESC").Limit(pageSize).Offset((page - 1) * pageSize).Find(&historyList).Error
|
|
|
|
if err != nil {
|
|
return Response{Code: 1, Message: err.Error()}
|
|
}
|
|
|
|
historyData := HistoryData{Page: page, PageSize: pageSize, Total: total, List: historyList}
|
|
|
|
return Response{Code: 0, Message: "success", Data: historyData}
|
|
}
|
|
|
|
func (a *App) Detect(filename string) Response {
|
|
db, err := a.GormDB()
|
|
if err != nil {
|
|
return Response{Code: 1, Message: err.Error()}
|
|
}
|
|
|
|
filePath := filepath.Join(publicImagePath, filename)
|
|
imgData, err := os.ReadFile(filePath)
|
|
if err != nil {
|
|
return Response{Code: 1, Message: "failed to read image: " + err.Error()}
|
|
}
|
|
|
|
input, err := preprocessImage(imgData)
|
|
if err != nil {
|
|
return Response{Code: 1, Message: "failed to preprocess: " + err.Error()}
|
|
}
|
|
|
|
// 推理
|
|
detectRet, confidence, err := runInference(input)
|
|
if err != nil {
|
|
return Response{Code: 1, Message: "model inference failed: " + err.Error()}
|
|
}
|
|
|
|
var breed Breed
|
|
err = db.Table("breeds_test").Where("code = ?", detectRet).First(&breed).Error
|
|
if err != nil {
|
|
return Response{Code: 1, Message: err.Error()}
|
|
}
|
|
|
|
now := int(time.Now().Unix())
|
|
|
|
one := HistoryItem{Img: filename, Breed: int(breed.Id), Date: now}
|
|
result := db.Table("history_test").Create(&one)
|
|
if result.Error != nil {
|
|
return Response{Code: 1, Message: result.Error.Error()}
|
|
}
|
|
|
|
detectData := DetectData{
|
|
Id: breed.Id,
|
|
Code: breed.Code,
|
|
Name: breed.Name,
|
|
Brief: breed.Brief,
|
|
ConfidenceLevel: confidence,
|
|
}
|
|
|
|
return Response{Code: 0, Message: "success", Data: detectData}
|
|
}
|
|
|
|
func (a *App) DeleteOneHistory(id uint) Response {
|
|
db, err := a.GormDB()
|
|
if err != nil {
|
|
return Response{Code: 1, Message: err.Error()}
|
|
}
|
|
|
|
result := db.Table("history_test").Delete(&HistoryItem{}, id)
|
|
if result.Error != nil {
|
|
return Response{Code: 1, Message: result.Error.Error()}
|
|
}
|
|
|
|
return Response{Code: 0, Message: "success"}
|
|
}
|
|
|
|
func (a *App) ClearHistory() Response {
|
|
db, err := a.GormDB()
|
|
if err != nil {
|
|
return Response{Code: 1, Message: err.Error()}
|
|
}
|
|
|
|
result := db.Exec("DELETE FROM history_test")
|
|
if result.Error != nil {
|
|
return Response{Code: 1, Message: result.Error.Error()}
|
|
}
|
|
|
|
if err := os.RemoveAll(publicImagePath); err != nil {
|
|
return Response{Code: 1, Message: fmt.Sprintf("failed to remove images: %v", err)}
|
|
}
|
|
if err := os.MkdirAll(publicImagePath, 0755); err != nil {
|
|
return Response{Code: 1, Message: fmt.Sprintf("failed to recreate images dir: %v", err)}
|
|
}
|
|
|
|
return Response{Code: 0, Message: "success"}
|
|
}
|