new train

backend/app.go

@@ -1,19 +1,39 @@
 package backend
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
 	"fmt"
+	"math"
 	"os"
 	"path/filepath"
 	"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 {
@@ -22,6 +42,113 @@ func NewApp() *App {
 
 func (a *App) Startup(ctx context.Context) {
 	a.ctx = ctx
+
+	// 设置 ONNX Runtime DLL 路径
+	// ort.SetSharedLibraryPath("onnxruntime.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) {
@@ -33,6 +160,7 @@ func (a *App) GormDB() (*gorm.DB, error) {
 }
 
 func (a *App) UploadImage(data []byte, filename string) Response {
+	println("UploadImage")
 	uploadsDir := publicImagePath
 	err := os.MkdirAll(uploadsDir, 0755)
 	if err != nil {
@@ -48,6 +176,8 @@ func (a *App) UploadImage(data []byte, filename string) Response {
 		return Response{Code: 1, Message: "failed", Data: err.Error()}
 	}
 
+	println("333")
+
 	return Response{Code: 0, Message: "success", Data: newFilename}
 }
 
@@ -94,14 +224,41 @@ func (a *App) GetHistory(page int, pageSize int) Response {
 	return Response{Code: 0, Message: "success", Data: historyData}
 }
 
-func (a *App) Detect(img string) Response {
+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()}
+	}
+
+	println("confidence", confidence)
+	println("detectRet", detectRet)
+
+	/**
+		结束
+	**/
+
 	var breed Breed
-	detectRet := "british_shorthair"
 	err = db.Table("breeds_test").Where("code = ?", detectRet).First(&breed).Error
 	if err != nil {
 		return Response{Code: 1, Message: err.Error()}
@@ -109,7 +266,7 @@ func (a *App) Detect(img string) Response {
 
 	now := int(time.Now().Unix())
 
-	one := HistoryItem{Img: img, Breed: int(breed.Id), Date: now}
+	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()}