-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathyolonasint8.go
107 lines (89 loc) · 2.54 KB
/
yolonasint8.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
96
97
98
99
100
101
102
103
104
105
106
107
package yolotriton
import (
"image"
"math"
triton "github.com/dev6699/yolotriton/grpc-client"
)
type YoloNASInt8 struct {
YoloTritonConfig
metadata struct {
xOffset float32
yOffset float32
scaleFactor float32
}
}
func NewYoloNASInt8(cfg YoloTritonConfig) Model {
return &YoloNASInt8{
YoloTritonConfig: cfg,
}
}
var _ Model = &YoloNAS{}
func (y *YoloNASInt8) GetConfig() YoloTritonConfig {
return y.YoloTritonConfig
}
func (y *YoloNASInt8) PreProcess(img image.Image, targetWidth uint, targetHeight uint) (*triton.InferTensorContents, error) {
height := img.Bounds().Dy()
width := img.Bounds().Dx()
scaleFactor := math.Min(float64(636)/float64(height), float64(636)/float64(width))
if scaleFactor != 1.0 {
newHeight := uint(math.Round(float64(height) * scaleFactor))
newWidth := uint(math.Round(float64(width) * scaleFactor))
img = resizeImage(img, newWidth, newHeight)
}
paddedImage, xOffset, yOffset := padImageToCenterWithGray(img, int(targetWidth), int(targetHeight), 114)
uint32Contents := imageToUint32Slice(paddedImage)
y.metadata.xOffset = float32(xOffset)
y.metadata.yOffset = float32(yOffset)
y.metadata.scaleFactor = float32(scaleFactor)
contents := &triton.InferTensorContents{
UintContents: uint32Contents,
}
return contents, nil
}
func (y *YoloNASInt8) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
numPreds, err := bytesToInt32Slice(rawOutputContents[0])
if err != nil {
return nil, err
}
predBoxes, err := bytesToFloat32Slice(rawOutputContents[1])
if err != nil {
return nil, err
}
predScores, err := bytesToFloat32Slice(rawOutputContents[2])
if err != nil {
return nil, err
}
predClasses, err := bytesToInt32Slice(rawOutputContents[3])
if err != nil {
return nil, err
}
boxes := []Box{}
detectedObjects := int(numPreds[0])
for index := 0; index < detectedObjects; index++ {
prob := predScores[index]
if prob < y.MinProbability {
continue
}
classID := predClasses[index]
label := y.Classes[classID]
idx := (index * 4)
x1raw := predBoxes[idx]
y1raw := predBoxes[idx+1]
x2raw := predBoxes[idx+2]
y2raw := predBoxes[idx+3]
scale := y.metadata.scaleFactor
x1 := (x1raw - y.metadata.xOffset) / scale
y1 := (y1raw - y.metadata.yOffset) / scale
x2 := (x2raw - y.metadata.xOffset) / scale
y2 := (y2raw - y.metadata.yOffset) / scale
boxes = append(boxes, Box{
X1: float64(x1),
Y1: float64(y1),
X2: float64(x2),
Y2: float64(y2),
Probability: float64(prob),
Class: label,
})
}
return boxes, nil
}