sugarBecomer
121 lines
2.5 KiB
Go
121 lines
2.5 KiB
Go
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)
|
|
}
|