Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| bc2bb25578 | |||
| b37fc5cb24 | |||
| 614dd9917f | |||
| 00bc3432bd |
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -384,18 +384,14 @@ func (c *CpdailyAutoSignTask) ParseFormJson(ctx context.Context, rawJson map[str
|
||||
}
|
||||
list = append(list, dto)
|
||||
|
||||
var current uint = 1
|
||||
|
||||
if isPhoto := rows["isPhoto"]; isPhoto != nil && isPhoto.Int() == 1 {
|
||||
dto = &formtype.CpTaskFormItem{
|
||||
Title: "签到照片",
|
||||
ColName: "signImg",
|
||||
Type: formtype.ItemTypeImage,
|
||||
Required: true,
|
||||
Number: current,
|
||||
}
|
||||
list = append(list, dto)
|
||||
current++
|
||||
}
|
||||
|
||||
// 是否范围外签到
|
||||
@@ -404,7 +400,6 @@ func (c *CpdailyAutoSignTask) ParseFormJson(ctx context.Context, rawJson map[str
|
||||
ColName: "isMalposition",
|
||||
Type: formtype.ItemTypeSingle,
|
||||
Required: true,
|
||||
Number: current,
|
||||
Options: []formtype.CpTaskFormItemOption{
|
||||
{
|
||||
Content: "是",
|
||||
@@ -418,7 +413,6 @@ func (c *CpdailyAutoSignTask) ParseFormJson(ctx context.Context, rawJson map[str
|
||||
},
|
||||
}
|
||||
list = append(list, dto)
|
||||
current++
|
||||
|
||||
// 范围外签到原因
|
||||
dto = &formtype.CpTaskFormItem{
|
||||
@@ -427,16 +421,15 @@ func (c *CpdailyAutoSignTask) ParseFormJson(ctx context.Context, rawJson map[str
|
||||
Type: formtype.ItemTypeInput,
|
||||
Required: true,
|
||||
LogicWid: "abnormalReason",
|
||||
Number: current,
|
||||
}
|
||||
list = append(list, dto)
|
||||
current++
|
||||
|
||||
current := uint(1)
|
||||
// 处理额外字段
|
||||
if isNeedExtra := rows["isNeedExtra"]; isNeedExtra != nil && isNeedExtra.Int() == 1 {
|
||||
extraField := rows["extraField"].Array()
|
||||
for i, field := range extraField {
|
||||
fieldObj := field.(map[string]*gvar.Var)
|
||||
for _, field := range extraField {
|
||||
fieldObj := gvar.New(field).MapStrVar()
|
||||
title := fieldObj["title"].String()
|
||||
isRequired := fieldObj["isRequired"].Bool()
|
||||
hasOtherItems := fieldObj["hasOtherItems"].Bool()
|
||||
@@ -445,16 +438,17 @@ func (c *CpdailyAutoSignTask) ParseFormJson(ctx context.Context, rawJson map[str
|
||||
Title: title,
|
||||
ColName: title,
|
||||
Type: formtype.ItemTypeSingle,
|
||||
Number: uint(i),
|
||||
Number: current,
|
||||
Required: isRequired,
|
||||
IsOther: hasOtherItems,
|
||||
}
|
||||
current++
|
||||
|
||||
// 处理选项
|
||||
extraFieldItems := fieldObj["extraFieldItems"].Array()
|
||||
var options []formtype.CpTaskFormItemOption
|
||||
for _, item := range extraFieldItems {
|
||||
itemObj := item.(map[string]*gvar.Var)
|
||||
itemObj := gvar.New(item).MapStrVar()
|
||||
option := formtype.CpTaskFormItemOption{
|
||||
Content: itemObj["content"].String(),
|
||||
IsOtherItems: itemObj["isOtherItems"].Bool(),
|
||||
@@ -577,8 +571,9 @@ func (c *CpdailyAutoSignTask) FillCustomForm(ctx context.Context, form map[strin
|
||||
var extraFieldItemValues []map[string]*gvar.Var
|
||||
|
||||
for i, field := range extraField {
|
||||
current := uint(i + 1)
|
||||
// 检查数字是否在范围内
|
||||
found := slices.Contains(numbers, uint(i))
|
||||
found := slices.Contains(numbers, current)
|
||||
if !found {
|
||||
utils.OkOrPanic(errors.NewBusinessError(500, "非法的额外信息数据"), oops.WithContext(ctx).Wrap)
|
||||
}
|
||||
@@ -586,7 +581,7 @@ func (c *CpdailyAutoSignTask) FillCustomForm(ctx context.Context, form map[strin
|
||||
// 找到对应的参数
|
||||
var targetParam *formtype.CpTaskFormItem
|
||||
for _, v := range paramList {
|
||||
if v.Number == uint(i) {
|
||||
if v.Number == current {
|
||||
targetParam = v
|
||||
break
|
||||
}
|
||||
|
||||
@@ -463,18 +463,15 @@ func (c *CpdailySleepCheckTask) ParseFormJson(ctx context.Context, rawJson map[s
|
||||
}
|
||||
list = append(list, dto)
|
||||
|
||||
var current uint = 1
|
||||
|
||||
if isPhoto := rows["isPhoto"]; isPhoto != nil && isPhoto.Int() == 1 {
|
||||
dto = &formtype.CpTaskFormItem{
|
||||
Title: "查寝照片",
|
||||
ColName: "signImg",
|
||||
Type: formtype.ItemTypeImage,
|
||||
Required: true,
|
||||
Number: current,
|
||||
}
|
||||
list = append(list, dto)
|
||||
current++
|
||||
|
||||
}
|
||||
|
||||
// 是否范围外签到
|
||||
@@ -483,7 +480,6 @@ func (c *CpdailySleepCheckTask) ParseFormJson(ctx context.Context, rawJson map[s
|
||||
ColName: "isMalposition",
|
||||
Type: formtype.ItemTypeSingle,
|
||||
Required: true,
|
||||
Number: current,
|
||||
Options: []formtype.CpTaskFormItemOption{
|
||||
{
|
||||
Content: "是",
|
||||
@@ -497,7 +493,6 @@ func (c *CpdailySleepCheckTask) ParseFormJson(ctx context.Context, rawJson map[s
|
||||
},
|
||||
}
|
||||
list = append(list, dto)
|
||||
current++
|
||||
|
||||
for key, item := range list {
|
||||
list[key].Show = item.LogicWid == ""
|
||||
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user