diff --git a/internal/application/classifier/classifier.go b/internal/application/classifier/classifier.go deleted file mode 100644 index 55586bb..0000000 --- a/internal/application/classifier/classifier.go +++ /dev/null @@ -1,225 +0,0 @@ -package classifier - -import ( - "bufio" - "fmt" - _ "image/jpeg" - _ "image/png" - "likedream-goframe/internal/framework/logger" - "os" - "path/filepath" - "runtime" - "strings" - - ort "github.com/yalue/onnxruntime_go" - "go.uber.org/zap" -) - -type ImageClassifier struct { - session *ort.AdvancedSession - inputShape []int64 - classes []string - inputTensor *ort.Tensor[float32] - outputTensor *ort.Tensor[float32] -} - -// NewImageClassifier 创建新的图片分类器 -func NewImageClassifier(modelPath, classesPath string) (*ImageClassifier, error) { - // 初始化ONNX Runtime - exePath := strings.ToLower(os.Args[0]) - // 获取绝对路径 - absPath, err := filepath.Abs(exePath) - dir := filepath.Dir(absPath) - systemOS := runtime.GOOS - switch systemOS { - case "windows": - dir = filepath.Join(dir, "onnxruntime.dll") - break - case "darwin": - dir = filepath.Join(dir, "onnxruntime.dylib") - break - case "linux": - dir = filepath.Join(dir, "onnxruntime.so") - break - } - ort.SetSharedLibraryPath(dir) - err = ort.InitializeEnvironment() - if err != nil { - logger.Logger.Error("初始化onnx:", zap.String("绝对路径", absPath), zap.String("父文件夹", dir)) - return nil, fmt.Errorf("初始化ONNX Runtime失败: %v", err) - } - - // 加载类别名称 - classes, err := loadClasses(classesPath) - if err != nil { - return nil, fmt.Errorf("加载类别失败: %v", err) - } - - // 创建session选项 - options, err := ort.NewSessionOptions() - if err != nil { - return nil, fmt.Errorf("创建session选项失败: %v", err) - } - defer options.Destroy() - - // 创建输入输出tensor - inputShape := ort.NewShape(1, 3, 224, 224) - outputShape := ort.NewShape(1, int64(len(classes))) - - inputTensor, err := ort.NewEmptyTensor[float32](inputShape) - if err != nil { - return nil, fmt.Errorf("创建输入tensor失败: %v", err) - } - - outputTensor, err := ort.NewEmptyTensor[float32](outputShape) - if err != nil { - inputTensor.Destroy() - return nil, fmt.Errorf("创建输出tensor失败: %v", err) - } - - // 加载模型 - session, err := ort.NewAdvancedSession(modelPath, - []string{"input"}, []string{"output"}, - []ort.ArbitraryTensor{inputTensor}, []ort.ArbitraryTensor{outputTensor}, - options) - if err != nil { - inputTensor.Destroy() - outputTensor.Destroy() - return nil, fmt.Errorf("加载模型失败: %v", err) - } - - return &ImageClassifier{ - session: session, - inputShape: []int64{1, 3, 224, 224}, - classes: classes, - inputTensor: inputTensor, - outputTensor: outputTensor, - }, nil -} - -// loadClasses 从文件加载类别名称 -func loadClasses(path string) ([]string, error) { - file, err := os.Open(path) - if err != nil { - return nil, err - } - defer file.Close() - - var classes []string - scanner := bufio.NewScanner(file) - for scanner.Scan() { - classes = append(classes, scanner.Text()) - } - - if err := scanner.Err(); err != nil { - return nil, err - } - - return classes, nil -} - -// Predict 预测图片类别 -func (ic *ImageClassifier) Predict(imgPath string) (string, float32, error) { - // 加载并预处理图片 - inputData, err := ic.preprocessImage(imgPath) - if err != nil { - return "", 0, fmt.Errorf("预处理图片失败: %v", err) - } - - // 设置输入数据 - copy(ic.inputTensor.GetData(), inputData) - - // 运行推理 - err = ic.session.Run() - if err != nil { - return "", 0, fmt.Errorf("推理失败: %v", err) - } - - // 获取输出数据 - outputData := ic.outputTensor.GetData() - - // 找到最大概率的类别 - maxIdx := 0 - maxProb := outputData[0] - for i, prob := range outputData { - if prob > maxProb { - maxProb = prob - maxIdx = i - } - } - - // 应用softmax获取概率 - expSum := float32(0) - for _, v := range outputData { - expSum += exp(v) - } - probability := exp(outputData[maxIdx]) / expSum - - return ic.classes[maxIdx], probability, nil -} - -// PredictFromBytes 从字节数据预测图片类别 -func (ic *ImageClassifier) PredictFromBytes(imgData []byte) (string, float32, error) { - // 预处理图片 - inputData, err := ic.preprocessImageBytes(imgData) - if err != nil { - return "", 0, fmt.Errorf("预处理图片失败: %v", err) - } - - // 设置输入数据 - copy(ic.inputTensor.GetData(), inputData) - - // 运行推理 - err = ic.session.Run() - if err != nil { - return "", 0, fmt.Errorf("推理失败: %v", err) - } - - // 获取输出数据 - outputData := ic.outputTensor.GetData() - - // 找到最大概率的类别 - maxIdx := 0 - maxProb := outputData[0] - for i, prob := range outputData { - if prob > maxProb { - maxProb = prob - maxIdx = i - } - } - - // 应用softmax获取概率 - expSum := float32(0) - for _, v := range outputData { - expSum += exp(v) - } - probability := exp(outputData[maxIdx]) / expSum - - return ic.classes[maxIdx], probability, nil -} - -// Close 关闭分类器 -func (ic *ImageClassifier) Close() error { - if ic.inputTensor != nil { - ic.inputTensor.Destroy() - } - if ic.outputTensor != nil { - ic.outputTensor.Destroy() - } - if ic.session != nil { - ic.session.Destroy() - } - return nil -} - -// exp 计算e的x次方(简单实现) -func exp(x float32) float32 { - // 使用泰勒级数近似 - result := float32(1.0) - term := float32(1.0) - for i := 1; i < 20; i++ { - term *= x / float32(i) - result += term - } - return result -} diff --git a/internal/application/classifier/preprocess.go b/internal/application/classifier/preprocess.go deleted file mode 100644 index 39fd49e..0000000 --- a/internal/application/classifier/preprocess.go +++ /dev/null @@ -1,120 +0,0 @@ -package classifier - -import ( - "bytes" - "image" - "os" -) - -// preprocessImage 预处理图片文件 -func (ic *ImageClassifier) preprocessImage(imgPath string) ([]float32, error) { - file, err := os.Open(imgPath) - if err != nil { - return nil, err - } - defer file.Close() - - img, _, err := image.Decode(file) - if err != nil { - return nil, err - } - - return ic.processImage(img) -} - -// preprocessImageBytes 预处理图片字节数据 -func (ic *ImageClassifier) preprocessImageBytes(imgData []byte) ([]float32, error) { - img, _, err := image.Decode(bytes.NewReader(imgData)) - if err != nil { - return nil, err - } - - return ic.processImage(img) -} - -// processImage 处理图片为模型输入格式 -func (ic *ImageClassifier) processImage(img image.Image) ([]float32, error) { - // 调整图片大小到224x224 - resized := resizeImage(img, 224, 224) - - // 转换为float32数组,并进行归一化 - // ImageNet的均值和标准差 - mean := []float32{0.485, 0.456, 0.406} - std := []float32{0.229, 0.224, 0.225} - - data := make([]float32, 3*224*224) - - for y := 0; y < 224; y++ { - for x := 0; x < 224; x++ { - r, g, b, _ := resized.At(x, y).RGBA() - - // 转换为0-1范围 - rNorm := float32(r) / 65535.0 - gNorm := float32(g) / 65535.0 - bNorm := float32(b) / 65535.0 - - // 应用归一化 - idx := y*224 + x - data[idx] = (rNorm - mean[0]) / std[0] // R通道 - data[224*224+idx] = (gNorm - mean[1]) / std[1] // G通道 - data[2*224*224+idx] = (bNorm - mean[2]) / std[2] // B通道 - } - } - - return data, nil -} - -// resizeImage 调整图片大小 -func resizeImage(img image.Image, width, height int) image.Image { - bounds := img.Bounds() - srcW := bounds.Dx() - srcH := bounds.Dy() - - dst := image.NewRGBA(image.Rect(0, 0, width, height)) - - for y := 0; y < height; y++ { - for x := 0; x < width; x++ { - // 简单的最近邻插值 - srcX := x * srcW / width - srcY := y * srcH / height - - if srcX >= srcW { - srcX = srcW - 1 - } - if srcY >= srcH { - srcY = srcH - 1 - } - - dst.Set(x, y, img.At(srcX, srcY)) - } - } - - return dst -} - -// centerCrop 中心裁剪图片 -func centerCrop(img image.Image, size int) image.Image { - bounds := img.Bounds() - width := bounds.Dx() - height := bounds.Dy() - - var x0, y0 int - if width > height { - x0 = (width - height) / 2 - y0 = 0 - width = height - } else { - x0 = 0 - y0 = (height - width) / 2 - height = width - } - - cropped := image.NewRGBA(image.Rect(0, 0, width, height)) - for y := 0; y < height; y++ { - for x := 0; x < width; x++ { - cropped.Set(x, y, img.At(x0+x, y0+y)) - } - } - - return resizeImage(cropped, size, size) -} diff --git a/internal/application/classifier/resolver.go b/internal/application/classifier/resolver.go deleted file mode 100644 index 772bb2b..0000000 --- a/internal/application/classifier/resolver.go +++ /dev/null @@ -1,17 +0,0 @@ -package classifier - -import "likedream-goframe/internal/utils" - -var imageClassifier, _ = NewImageClassifier("static/model.onnx", "static/classes.txt") - -// 图片数据 & 类别名称 -// return 是否匹配 -func ClassifyImage(imgData []byte, className string) bool { - if imageClassifier == nil { - c, err := NewImageClassifier("static/model.onnx", "static/classes.txt") - utils.OkOrPanic(err) - imageClassifier = c - } - name, _, _ := imageClassifier.PredictFromBytes(imgData) - return className == name -} diff --git a/internal/application/recognizer/init.go b/internal/application/recognizer/init.go index 4a0d591..8f4c77e 100644 --- a/internal/application/recognizer/init.go +++ b/internal/application/recognizer/init.go @@ -43,15 +43,15 @@ func getOnnxRuntimeName() (onnxRuntimeFileName string, onnxRuntimeFileBytes []by switch systemOS { case "windows": onnxRuntimeFileName = "onnxruntime.dll" - onnxRuntimeFileBytes, _ = fs.ReadFile("static/" + onnxRuntimeFileName) + onnxRuntimeFileBytes, _ = fs.ReadFile("static/onnxruntime.dll") return case "darwin": onnxRuntimeFileName = "onnxruntime.dylib" - onnxRuntimeFileBytes, _ = fs.ReadFile("static/" + onnxRuntimeFileName) + onnxRuntimeFileBytes, _ = fs.ReadFile("static/libonnxruntime.1.23.2.dylib") return case "linux": onnxRuntimeFileName = "onnxruntime.so" - onnxRuntimeFileBytes, _ = fs.ReadFile("static/" + onnxRuntimeFileName) + onnxRuntimeFileBytes, _ = fs.ReadFile("static/libonnxruntime.so.1.23.2") return } logger.Logger.Panic("暂不支持的操作系统:", zap.String("系统名称", systemOS))