likedream/internal/application/recognizer/classifier/classifier.go

package classifier

import (
	"bufio"
	"fmt"
	_ "image/jpeg"
	_ "image/png"
	"os"

	ort "github.com/yalue/onnxruntime_go"
)

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) {
	// 加载类别名称
	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
}