18 lines
501 B
Go
18 lines
501 B
Go
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
|
|
}
|