sugarBecomer
60 lines
1.5 KiB
Go
60 lines
1.5 KiB
Go
package recognizer
|
||
|
||
import (
|
||
"embed"
|
||
_ "embed"
|
||
"likedream-goframe/internal/framework/logger"
|
||
"os"
|
||
"path/filepath"
|
||
"runtime"
|
||
"strings"
|
||
|
||
ort "github.com/yalue/onnxruntime_go"
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
//go:embed static/*
|
||
var fs embed.FS
|
||
|
||
// OnnxIsInit 初始化onnx配置
|
||
var OnnxIsInit = false
|
||
|
||
func InitOnnx() {
|
||
name, bytes := getOnnxRuntimeName()
|
||
// 检测文件是否存在
|
||
exePath := strings.ToLower(os.Args[0])
|
||
// 获取绝对路径
|
||
absPath, _ := filepath.Abs(exePath)
|
||
filePath := filepath.Join(filepath.Dir(absPath), name)
|
||
if _, err := os.Stat(filePath); os.IsNotExist(err) {
|
||
os.WriteFile(filePath, bytes, 0644)
|
||
}
|
||
ort.SetSharedLibraryPath(filePath)
|
||
err := ort.InitializeEnvironment()
|
||
if err != nil {
|
||
logger.Logger.Panic("初始化onnx：", zap.String("绝对路径", filePath))
|
||
return
|
||
}
|
||
OnnxIsInit = true
|
||
}
|
||
|
||
func getOnnxRuntimeName() (onnxRuntimeFileName string, onnxRuntimeFileBytes []byte) {
|
||
systemOS := runtime.GOOS
|
||
switch systemOS {
|
||
case "windows":
|
||
onnxRuntimeFileName = "onnxruntime.dll"
|
||
onnxRuntimeFileBytes, _ = fs.ReadFile("static/onnxruntime.dll")
|
||
return
|
||
case "darwin":
|
||
onnxRuntimeFileName = "onnxruntime.dylib"
|
||
onnxRuntimeFileBytes, _ = fs.ReadFile("static/libonnxruntime.1.23.2.dylib")
|
||
return
|
||
case "linux":
|
||
onnxRuntimeFileName = "onnxruntime.so"
|
||
onnxRuntimeFileBytes, _ = fs.ReadFile("static/libonnxruntime.so.1.23.2")
|
||
return
|
||
}
|
||
logger.Logger.Panic("暂不支持的操作系统：", zap.String("系统名称", systemOS))
|
||
return "", nil
|
||
}
|