-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathyolov8.go
95 lines (74 loc) · 1.99 KB
/
yolov8.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
package yolotriton
import (
"image"
triton "github.com/dev6699/yolotriton/grpc-client"
)
type YoloV8 struct {
YoloTritonConfig
metadata struct {
scaleFactorW float32
scaleFactorH float32
}
}
func NewYoloV8(cfg YoloTritonConfig) Model {
return &YoloV8{
YoloTritonConfig: cfg,
}
}
var _ Model = &YoloV8{}
func (y *YoloV8) GetConfig() YoloTritonConfig {
return y.YoloTritonConfig
}
func (y *YoloV8) PreProcess(img image.Image, targetWidth uint, targetHeight uint) (*triton.InferTensorContents, error) {
width := img.Bounds().Dx()
height := img.Bounds().Dy()
preprocessedImg := resizeImage(img, targetWidth, targetHeight)
fp32Contents := imageToFloat32Slice(preprocessedImg)
y.metadata.scaleFactorW = float32(width) / float32(targetWidth)
y.metadata.scaleFactorH = float32(height) / float32(targetHeight)
contents := &triton.InferTensorContents{
Fp32Contents: fp32Contents,
}
return contents, nil
}
func (y *YoloV8) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
output, err := bytesToFloat32Slice(rawOutputContents[0])
if err != nil {
return nil, err
}
numObjects := y.NumObjects
numClasses := y.NumClasses
boxes := []Box{}
for index := 0; index < numObjects; index++ {
classID := 0
prob := float32(0.0)
for col := 0; col < numClasses; col++ {
p := output[numObjects*(col+4)+index]
if p > prob {
prob = p
classID = col
}
}
if prob < y.MinProbability {
continue
}
label := y.Classes[classID]
x1raw := output[index]
y1raw := output[numObjects+index]
w := output[2*numObjects+index]
h := output[3*numObjects+index]
x1 := (x1raw - w/2) * y.metadata.scaleFactorW
y1 := (y1raw - h/2) * y.metadata.scaleFactorH
x2 := (x1raw + w/2) * y.metadata.scaleFactorW
y2 := (y1raw + h/2) * y.metadata.scaleFactorH
boxes = append(boxes, Box{
X1: float64(x1),
Y1: float64(y1),
X2: float64(x2),
Y2: float64(y2),
Probability: float64(prob),
Class: label,
})
}
return boxes, nil
}