apps/yolotriton
1
0
Fork 0
mirror of https://github.com/dev6699/yolotriton.git synced 2025-09-26 19:51:13 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: added support for YOLO-NAS INT8

Browse Source
This commit is contained in:
kweijack
2023-09-15 04:36:18 +00:00
parent 8ee97b8bca
commit 1b32cbc3b3
9 changed files with 372 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
cmd/main.go
View File

@@ -5,25 +5,30 @@ import (
"fmt"
"log"
"strings"
"time"
"github.com/dev6699/yolotriton"
)
type Flags struct {
ModelName string
ModelVersion string
ModelType string
URL string
Image string
ModelName string
ModelVersion string
ModelType string
URL string
Image string
Benchmark bool
BenchmarkCount int
}
func parseFlags() Flags {
var flags Flags
flag.StringVar(&flags.ModelName, "m", "yolonas", "Name of model being served (Required)")
flag.StringVar(&flags.ModelVersion, "x", "", "Version of model. Default: Latest Version")
flag.StringVar(&flags.ModelType, "t", "yolonas", "Type of model. Available options: [yolonas, yolov8]")
flag.StringVar(&flags.ModelType, "t", "yolonas", "Type of model. Available options: [yolonas, yolonasint8, yolov8]")
flag.StringVar(&flags.URL, "u", "tritonserver:8001", "Inference Server URL.")
flag.StringVar(&flags.Image, "i", "images/1.jpg", "Inference Image.")
flag.BoolVar(&flags.Benchmark, "b", false, "Run benchmark.")
flag.IntVar(&flags.BenchmarkCount, "n", 1, "Number of benchmark run.")
flag.Parse()
return flags
}
@@ -38,11 +43,13 @@ func main() {
model = yolotriton.NewYoloV8(FLAGS.ModelName, FLAGS.ModelVersion)
case yolotriton.ModelTypeYoloNAS:
model = yolotriton.NewYoloNAS(FLAGS.ModelName, FLAGS.ModelVersion)
case yolotriton.ModelTypeYoloNASInt8:
model = yolotriton.NewYoloNASInt8(FLAGS.ModelName, FLAGS.ModelVersion)
default:
log.Fatalf("Unsupported model: %s. Available options: [yolonas, yolov8]", FLAGS.ModelType)
log.Fatalf("Unsupported model: %s. Available options: [yolonas, yolonasint8, yolov8]", FLAGS.ModelType)
}
ygt, err := yolotriton.New(FLAGS.URL, model)
yt, err := yolotriton.New(FLAGS.URL, model)
if err != nil {
log.Fatal(err)
}
@@ -52,31 +59,47 @@ func main() {
log.Fatalf("Failed to preprocess image: %v", err)
}
results, err := ygt.Infer(img)
if err != nil {
log.Fatal(err)
loop := 1
if FLAGS.Benchmark {
loop = FLAGS.BenchmarkCount
}
for i, r := range results {
fmt.Println("prediction: ", i)
fmt.Println("class: ", r.Class)
fmt.Printf("confidence: %.2f\n", r.Probability)
fmt.Println("bboxes: [", int(r.X1), int(r.Y1), int(r.X2), int(r.Y2), "]")
fmt.Println("---------------------")
start := time.Now()
for i := 0; i < loop; i++ {
now := time.Now()
results, err := yt.Infer(img)
if err != nil {
log.Fatal(err)
}
fmt.Printf("%d. processing time: %s\n", i+1, time.Since(now))
if FLAGS.Benchmark {
continue
}
for i, r := range results {
fmt.Println("prediction: ", i)
fmt.Println("class: ", r.Class)
fmt.Printf("confidence: %.2f\n", r.Probability)
fmt.Println("bboxes: [", int(r.X1), int(r.Y1), int(r.X2), int(r.Y2), "]")
fmt.Println("---------------------")
}
out, err := yolotriton.DrawBoundingBoxes(
img,
results,
int(float64(img.Bounds().Dx())*0.005),
float64(img.Bounds().Dx())*0.02,
)
if err != nil {
log.Fatal(err)
}
err = yolotriton.SaveImage(out, fmt.Sprintf("%s_%s_out.jpg", strings.Split(FLAGS.Image, ".")[0], FLAGS.ModelName))
if err != nil {
log.Fatal(err)
}
}
out, err := yolotriton.DrawBoundingBoxes(
img,
results,
int(float64(img.Bounds().Dx())*0.005),
float64(img.Bounds().Dx())*0.02,
)
if err != nil {
log.Fatal(err)
}
err = yolotriton.SaveImage(out, fmt.Sprintf("%s_%s_out.jpg", strings.Split(FLAGS.Image, ".")[0], FLAGS.ModelName))
if err != nil {
log.Fatal(err)
if FLAGS.Benchmark {
fmt.Println("Avg processing time:", time.Since(start)/time.Duration(FLAGS.BenchmarkCount))
}
}

0
model_repository/yolonasint8/1/.gitkeep Normal file
View File

2
model_repository/yolonasint8/config.pbtxt Normal file
View File

@@ -0,0 +1,2 @@
name: "yolonasint8"
platform: "tensorrt_plan"

23
postprocess.go
View File

@@ -30,6 +30,29 @@ func bytesToFloat32Slice(data []byte) ([]float32, error) {
return t, nil
}
func bytesToInt32Slice(data []byte) ([]int32, error) {
t := []int32{}
// Create a buffer from the input data
buffer := bytes.NewReader(data)
for {
// Read the binary data from the buffer
var binaryValue uint32
err := binary.Read(buffer, binary.LittleEndian, &binaryValue)
if err != nil {
if err == io.EOF {
break
}
return nil, err
}
t = append(t, int32(binaryValue))
}
return t, nil
}
type Box struct {
X1 float64
Y1 float64

50
preprocess.go
View File

@@ -3,6 +3,8 @@ package yolotriton
import (
"image"
"image/color"
"image/draw"
"math"
"github.com/nfnt/resize"
)
@@ -43,3 +45,51 @@ func imageToFloat32Slice(img image.Image) []float32 {
return inputContents
}
func imageToUint32Slice(img image.Image) []uint32 {
bounds := img.Bounds()
width, height := bounds.Max.X, bounds.Max.Y
inputContents := make([]uint32, width*height*3)
idx := 0
offset := (height * width)
for y := 0; y < height; y++ {
for x := 0; x < width; x++ {
pixel := img.At(x, y)
r, g, b, _ := pixelRGBA(pixel)
inputContents[idx] = r
inputContents[offset+idx] = g
inputContents[2*offset+idx] = b
idx++
}
}
return inputContents
}
func padImageToCenterWithGray(originalImage image.Image, targetWidth, targetHeight int, grayValue uint8) (image.Image, int, int) {
// Calculate the dimensions of the original image
originalWidth := originalImage.Bounds().Dx()
originalHeight := originalImage.Bounds().Dy()
// Calculate the padding dimensions
padWidth := targetWidth - originalWidth
padHeight := targetHeight - originalHeight
// Create a new RGBA image with the desired dimensions and fill it with gray
paddedImage := image.NewRGBA(image.Rect(0, 0, targetWidth, targetHeight))
grayColor := color.RGBA{grayValue, grayValue, grayValue, 255}
draw.Draw(paddedImage, paddedImage.Bounds(), &image.Uniform{grayColor}, image.Point{}, draw.Src)
// Calculate the position to paste the original image in the center
xOffset := int(math.Floor(float64(padWidth) / 2))
yOffset := int(math.Floor(float64(padHeight) / 2))
// Paste the original image onto the padded image
pasteRect := image.Rect(xOffset, yOffset, xOffset+originalWidth, yOffset+originalHeight)
draw.Draw(paddedImage, pasteRect, originalImage, image.Point{}, draw.Over)
return paddedImage, xOffset, yOffset
}

138
yolo.go
View File

@@ -13,19 +13,20 @@ import (
type ModelType string
const (
ModelTypeYoloV8 ModelType = "yolov8"
ModelTypeYoloNAS ModelType = "yolonas"
ModelTypeYoloV8 ModelType = "yolov8"
ModelTypeYoloNAS ModelType = "yolonas"
ModelTypeYoloNASInt8 ModelType = "yolonasint8"
)
type Model interface {
GetConfig() YoloTritonConfig
PreProcess(img image.Image, targetWidth uint, targetHeight uint) ([]float32, error)
PreProcess(img image.Image, targetWidth uint, targetHeight uint) (*triton.InferTensorContents, error)
PostProcess(rawOutputContents [][]byte) ([]Box, error)
GetClass(index int) string
}
type YoloTritonConfig struct {
BatchSize int
NumChannels int
NumClasses int
NumObjects int
ModelName string
ModelVersion string
@@ -39,17 +40,25 @@ func New(url string, model Model) (*YoloTriton, error) {
return nil, err
}
cfg := model.GetConfig()
modelMetadata, err := newModelMetadata(conn, cfg.ModelName, cfg.ModelVersion)
if err != nil {
return nil, err
}
return &YoloTriton{
conn: conn,
model: model,
cfg: model.GetConfig(),
conn: conn,
model: model,
cfg: cfg,
modelMetadata: modelMetadata,
}, nil
}
type YoloTriton struct {
cfg YoloTritonConfig
conn *grpc.ClientConn
model Model
model Model
cfg YoloTritonConfig
conn *grpc.ClientConn
modelMetadata *modelMetadata
}
func (y *YoloTriton) Close() error {
@@ -58,51 +67,14 @@ func (y *YoloTriton) Close() error {
func (y *YoloTriton) Infer(img image.Image) ([]Box, error) {
inputs, err := y.model.PreProcess(img, y.modelMetadata.inputWidth(), y.modelMetadata.inputHeight())
if err != nil {
return nil, err
}
modelInferRequest := y.modelMetadata.formInferRequest(inputs)
client := triton.NewGRPCInferenceServiceClient(y.conn)
metaResponse, err := ModelMetadataRequest(client, y.cfg.ModelName, y.cfg.ModelVersion)
if err != nil {
return nil, err
}
modelInferRequest := &triton.ModelInferRequest{
ModelName: y.cfg.ModelName,
ModelVersion: y.cfg.ModelVersion,
}
input := metaResponse.Inputs[0]
if input.Shape[0] == -1 {
input.Shape[0] = 1
}
inputWidth := input.Shape[2]
inputHeight := input.Shape[3]
fp32Contents, err := y.model.PreProcess(img, uint(inputWidth), uint(inputHeight))
if err != nil {
return nil, err
}
modelInferRequest.Inputs = append(modelInferRequest.Inputs,
&triton.ModelInferRequest_InferInputTensor{
Name: input.Name,
Datatype: input.Datatype,
Shape: input.Shape,
Contents: &triton.InferTensorContents{
// Simply assume all are fp32
Fp32Contents: fp32Contents,
},
},
)
for _, o := range metaResponse.Outputs {
modelInferRequest.Outputs = append(modelInferRequest.Outputs,
&triton.ModelInferRequest_InferRequestedOutputTensor{
Name: o.Name,
},
)
}
inferResponse, err := ModelInferRequest(client, modelInferRequest)
if err != nil {
return nil, err
@@ -131,3 +103,59 @@ func (y *YoloTriton) Infer(img image.Image) ([]Box, error) {
return result, nil
}
type modelMetadata struct {
modelName string
modelVersion string
*triton.ModelMetadataResponse
}
func newModelMetadata(conn *grpc.ClientConn, modelName string, modelVersion string) (*modelMetadata, error) {
client := triton.NewGRPCInferenceServiceClient(conn)
metaResponse, err := ModelMetadataRequest(client, modelName, modelVersion)
if err != nil {
return nil, err
}
return &modelMetadata{
modelName: modelName,
modelVersion: modelVersion,
ModelMetadataResponse: metaResponse,
}, nil
}
func (m *modelMetadata) inputWidth() uint {
return uint(m.Inputs[0].Shape[2])
}
func (m *modelMetadata) inputHeight() uint {
return uint(m.Inputs[0].Shape[3])
}
func (m *modelMetadata) formInferRequest(contents *triton.InferTensorContents) *triton.ModelInferRequest {
input := m.Inputs[0]
if input.Shape[0] == -1 {
input.Shape[0] = 1
}
outputs := make([]*triton.ModelInferRequest_InferRequestedOutputTensor, len(m.Outputs))
for i, o := range m.Outputs {
outputs[i] = &triton.ModelInferRequest_InferRequestedOutputTensor{
Name: o.Name,
}
}
return &triton.ModelInferRequest{
ModelName: m.modelName,
ModelVersion: m.modelVersion,
Inputs: []*triton.ModelInferRequest_InferInputTensor{
{
Name: input.Name,
Datatype: input.Datatype,
Shape: input.Shape,
Contents: contents,
},
},
Outputs: outputs,
}
}

67
yolonas.go
View File

@@ -2,9 +2,9 @@ package yolotriton
import (
"image"
"image/color"
"image/draw"
"math"
triton "github.com/dev6699/yolotriton/grpc-client"
)
type YoloNAS struct {
@@ -19,8 +19,7 @@ type YoloNAS struct {
func NewYoloNAS(modelName string, modelVersion string) Model {
return &YoloNAS{
YoloTritonConfig: YoloTritonConfig{
BatchSize: 1,
NumChannels: 80,
NumClasses: 80,
NumObjects: 8400,
MinProbability: 0.5,
MaxIOU: 0.7,
@@ -36,7 +35,11 @@ func (y *YoloNAS) GetConfig() YoloTritonConfig {
return y.YoloTritonConfig
}
func (y *YoloNAS) PreProcess(img image.Image, targetWidth uint, targetHeight uint) ([]float32, error) {
func (y *YoloNAS) GetClass(index int) string {
return yoloClasses[index]
}
func (y *YoloNAS) PreProcess(img image.Image, targetWidth uint, targetHeight uint) (*triton.InferTensorContents, error) {
height := img.Bounds().Dy()
width := img.Bounds().Dx()
@@ -56,7 +59,10 @@ func (y *YoloNAS) PreProcess(img image.Image, targetWidth uint, targetHeight uin
y.metadata.yOffset = float32(yOffset)
y.metadata.scaleFactor = float32(scaleFactor)
return fp32Contents, nil
contents := &triton.InferTensorContents{
Fp32Contents: fp32Contents,
}
return contents, nil
}
func (y *YoloNAS) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
@@ -76,8 +82,8 @@ func (y *YoloNAS) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
classID := 0
prob := float32(0.0)
for col := 0; col < y.NumChannels; col++ {
p := predScores[index*y.NumChannels+(col)]
for col := 0; col < y.NumClasses; col++ {
p := predScores[index*y.NumClasses+(col)]
if p > prob {
prob = p
classID = col
@@ -88,18 +94,18 @@ func (y *YoloNAS) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
continue
}
label := yoloClasses[classID]
i := (index * 4)
xc := predBoxes[i]
yc := predBoxes[i+1]
w := predBoxes[i+2]
h := predBoxes[i+3]
label := y.GetClass(classID)
idx := (index * 4)
x1raw := predBoxes[idx]
y1raw := predBoxes[idx+1]
x2raw := predBoxes[idx+2]
y2raw := predBoxes[idx+3]
scale := y.metadata.scaleFactor
x1 := (xc - y.metadata.xOffset) / scale
y1 := (yc - y.metadata.yOffset) / scale
x2 := (w - y.metadata.xOffset) / scale
y2 := (h - y.metadata.yOffset) / scale
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),
@@ -113,28 +119,3 @@ func (y *YoloNAS) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
return boxes, nil
}
func padImageToCenterWithGray(originalImage image.Image, targetWidth, targetHeight int, grayValue uint8) (image.Image, int, int) {
// Calculate the dimensions of the original image
originalWidth := originalImage.Bounds().Dx()
originalHeight := originalImage.Bounds().Dy()
// Calculate the padding dimensions
padWidth := targetWidth - originalWidth
padHeight := targetHeight - originalHeight
// Create a new RGBA image with the desired dimensions and fill it with gray
paddedImage := image.NewRGBA(image.Rect(0, 0, targetWidth, targetHeight))
grayColor := color.RGBA{grayValue, grayValue, grayValue, 255}
draw.Draw(paddedImage, paddedImage.Bounds(), &image.Uniform{grayColor}, image.Point{}, draw.Src)
// Calculate the position to paste the original image in the center
xOffset := int(math.Floor(float64(padWidth) / 2))
yOffset := int(math.Floor(float64(padHeight) / 2))
// Paste the original image onto the padded image
pasteRect := image.Rect(xOffset, yOffset, xOffset+originalWidth, yOffset+originalHeight)
draw.Draw(paddedImage, pasteRect, originalImage, image.Point{}, draw.Over)
return paddedImage, xOffset, yOffset
}

116
yolonasint8.go Normal file
View File

@@ -0,0 +1,116 @@
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(modelName string, modelVersion string) Model {
return &YoloNASInt8{
YoloTritonConfig: YoloTritonConfig{
MinProbability: 0.5,
MaxIOU: 0.7,
ModelName: modelName,
ModelVersion: modelVersion,
},
}
}
var _ Model = &YoloNAS{}
func (y *YoloNASInt8) GetConfig() YoloTritonConfig {
return y.YoloTritonConfig
}
func (y *YoloNASInt8) GetClass(index int) string {
return yoloClasses[index]
}
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.GetClass(int(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
}

34
yolov8.go
View File

@@ -2,6 +2,8 @@ package yolotriton
import (
"image"
triton "github.com/dev6699/yolotriton/grpc-client"
)
type YoloV8 struct {
@@ -15,8 +17,7 @@ type YoloV8 struct {
func NewYoloV8(modelName string, modelVersion string) Model {
return &YoloV8{
YoloTritonConfig: YoloTritonConfig{
BatchSize: 1,
NumChannels: 84,
NumClasses: 80,
NumObjects: 8400,
MinProbability: 0.5,
MaxIOU: 0.7,
@@ -32,7 +33,11 @@ func (y *YoloV8) GetConfig() YoloTritonConfig {
return y.YoloTritonConfig
}
func (y *YoloV8) PreProcess(img image.Image, targetWidth uint, targetHeight uint) ([]float32, error) {
func (y *YoloV8) GetClass(index int) string {
return yoloClasses[index]
}
func (y *YoloV8) PreProcess(img image.Image, targetWidth uint, targetHeight uint) (*triton.InferTensorContents, error) {
width := img.Bounds().Dx()
height := img.Bounds().Dy()
@@ -43,7 +48,10 @@ func (y *YoloV8) PreProcess(img image.Image, targetWidth uint, targetHeight uint
y.metadata.scaleFactorW = float32(width) / float32(targetWidth)
y.metadata.scaleFactorH = float32(height) / float32(targetHeight)
return fp32Contents, nil
contents := &triton.InferTensorContents{
Fp32Contents: fp32Contents,
}
return contents, nil
}
func (y *YoloV8) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
@@ -53,7 +61,7 @@ func (y *YoloV8) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
}
numObjects := y.NumObjects
numChannels := y.NumChannels
numClasses := y.NumClasses
boxes := []Box{}
@@ -61,7 +69,7 @@ func (y *YoloV8) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
classID := 0
prob := float32(0.0)
for col := 0; col < numChannels-4; col++ {
for col := 0; col < numClasses; col++ {
p := output[numObjects*(col+4)+index]
if p > prob {
prob = p
@@ -73,16 +81,16 @@ func (y *YoloV8) PostProcess(rawOutputContents [][]byte) ([]Box, error) {
continue
}
label := yoloClasses[classID]
xc := output[index]
yc := output[numObjects+index]
label := y.GetClass(classID)
x1raw := output[index]
y1raw := output[numObjects+index]
w := output[2*numObjects+index]
h := output[3*numObjects+index]
x1 := (xc - w/2) * y.metadata.scaleFactorW
y1 := (yc - h/2) * y.metadata.scaleFactorH
x2 := (xc + w/2) * y.metadata.scaleFactorW
y2 := (yc + h/2) * y.metadata.scaleFactorH
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),