sugarBecomer
54 lines
1.4 KiB
Go
54 lines
1.4 KiB
Go
package classifier
|
|
|
|
import (
|
|
"embed"
|
|
_ "embed"
|
|
"likedream-goframe/internal/application/recognizer"
|
|
"likedream-goframe/internal/framework/logger"
|
|
"likedream-goframe/internal/utils"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
//go:embed static/*
|
|
var fs embed.FS
|
|
var imageClassifier *ImageClassifier
|
|
|
|
func init() {
|
|
if !recognizer.OnnxIsInit {
|
|
recognizer.InitOnnx()
|
|
}
|
|
modelPath, classesPath := initModel()
|
|
classifier, err := NewImageClassifier(modelPath, classesPath)
|
|
if err != nil {
|
|
logger.Logger.Panic("初始化ai模型失败", zap.Error(err))
|
|
}
|
|
imageClassifier = classifier
|
|
}
|
|
func initModel() (modelPath, classesPath string) {
|
|
// 检测文件是否存在
|
|
exePath := strings.ToLower(os.Args[0])
|
|
// 获取绝对路径
|
|
absPath, _ := filepath.Abs(exePath)
|
|
staticPath := filepath.Join(filepath.Dir(absPath), "static")
|
|
if _, err := os.Stat(staticPath); os.IsNotExist(err) {
|
|
os.Mkdir(staticPath, 0755)
|
|
}
|
|
modelPath = filepath.Join(staticPath, "model.onnx")
|
|
modelBytes, err := fs.ReadFile("static/model.onnx")
|
|
utils.OkOrPanic(err)
|
|
if _, err := os.Stat(modelPath); os.IsNotExist(err) {
|
|
os.WriteFile(modelPath, modelBytes, 0644)
|
|
}
|
|
classesPath = filepath.Join(staticPath, "classes.txt")
|
|
classesBytes, err := fs.ReadFile("static/classes.txt")
|
|
utils.OkOrPanic(err)
|
|
if _, err := os.Stat(classesPath); os.IsNotExist(err) {
|
|
os.WriteFile(classesPath, classesBytes, 0644)
|
|
}
|
|
return
|
|
}
|