feat: added support for YOLO-NAS INT8

cmd/main.go

@@ -5,25 +5,30 @@ import (
 	"fmt"
 	"log"
 	"strings"
+	"time"
 
 	"github.com/dev6699/yolotriton"
 )
 
 type Flags struct {
-	ModelName    string
+	ModelName      string
-	ModelVersion string
+	ModelVersion   string
-	ModelType    string
+	ModelType      string
-	URL          string
+	URL            string
-	Image        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)
+	loop := 1
-	if err != nil {
+	if FLAGS.Benchmark {
-		log.Fatal(err)
+		loop = FLAGS.BenchmarkCount
 	}
 
-	for i, r := range results {
+	start := time.Now()
-		fmt.Println("prediction: ", i)
+	for i := 0; i < loop; i++ {
-		fmt.Println("class: ", r.Class)
+		now := time.Now()
-		fmt.Printf("confidence: %.2f\n", r.Probability)
+		results, err := yt.Infer(img)
-		fmt.Println("bboxes: [", int(r.X1), int(r.Y1), int(r.X2), int(r.Y2), "]")
+		if err != nil {
-		fmt.Println("---------------------")
+			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(
+	if FLAGS.Benchmark {
-		img,
+		fmt.Println("Avg processing time:", time.Since(start)/time.Duration(FLAGS.BenchmarkCount))
-		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)
 	}
 }

model_repository/yolonasint8/config.pbtxt

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

postprocess.go

@@ -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

preprocess.go

@@ -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
+}

yolo.go

@@ -13,19 +13,20 @@ import (
 type ModelType string
 
 const (
-	ModelTypeYoloV8  ModelType = "yolov8"
+	ModelTypeYoloV8      ModelType = "yolov8"
-	ModelTypeYoloNAS ModelType = "yolonas"
+	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
+	NumClasses     int
-	NumChannels    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,
+		conn:          conn,
-		model: model,
+		model:         model,
-		cfg:   model.GetConfig(),
+		cfg:           cfg,
+		modelMetadata: modelMetadata,
 	}, nil
 }
 
 type YoloTriton struct {
-	cfg   YoloTritonConfig
+	model         Model
-	conn  *grpc.ClientConn
+	cfg           YoloTritonConfig
-	model Model
+	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,
+	}
+}

yolonas.go

@@ -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,
+			NumClasses:     80,
-			NumChannels:    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++ {
+		for col := 0; col < y.NumClasses; col++ {
-			p := predScores[index*y.NumChannels+(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]
+		label := y.GetClass(classID)
-		i := (index * 4)
+		idx := (index * 4)
-		xc := predBoxes[i]
+		x1raw := predBoxes[idx]
-		yc := predBoxes[i+1]
+		y1raw := predBoxes[idx+1]
-		w := predBoxes[i+2]
+		x2raw := predBoxes[idx+2]
-		h := predBoxes[i+3]
+		y2raw := predBoxes[idx+3]
 
 		scale := y.metadata.scaleFactor
-		x1 := (xc - y.metadata.xOffset) / scale
+		x1 := (x1raw - y.metadata.xOffset) / scale
-		y1 := (yc - y.metadata.yOffset) / scale
+		y1 := (y1raw - y.metadata.yOffset) / scale
-		x2 := (w - y.metadata.xOffset) / scale
+		x2 := (x2raw - y.metadata.xOffset) / scale
-		y2 := (h - y.metadata.yOffset) / 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
-}

yolonasint8.go

@@ -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
+}

yolov8.go

@@ -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,
+			NumClasses:     80,
-			NumChannels:    84,
 			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]
+		label := y.GetClass(classID)
-		xc := output[index]
+		x1raw := output[index]
-		yc := output[numObjects+index]
+		y1raw := output[numObjects+index]
 		w := output[2*numObjects+index]
 		h := output[3*numObjects+index]
 
-		x1 := (xc - w/2) * y.metadata.scaleFactorW
+		x1 := (x1raw - w/2) * y.metadata.scaleFactorW
-		y1 := (yc - h/2) * y.metadata.scaleFactorH
+		y1 := (y1raw - h/2) * y.metadata.scaleFactorH
-		x2 := (xc + w/2) * y.metadata.scaleFactorW
+		x2 := (x1raw + w/2) * y.metadata.scaleFactorW
-		y2 := (yc + h/2) * y.metadata.scaleFactorH
+		y2 := (y1raw + h/2) * y.metadata.scaleFactorH
 
 		boxes = append(boxes, Box{
 			X1:          float64(x1),