From c1b5cfdf2d9ac67d8f9b8553264cbcb388396fc9 Mon Sep 17 00:00:00 2001
From: lindsayshuo <932695342@qq.com>
Date: Thu, 14 Mar 2024 13:53:42 +0800
Subject: [PATCH] Update yolov8_cls_trt.py (#1455)

batch image data
---
 yolov8/yolov8_cls_trt.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/yolov8/yolov8_cls_trt.py b/yolov8/yolov8_cls_trt.py
index 98c0846..ad75709 100644
--- a/yolov8/yolov8_cls_trt.py
+++ b/yolov8/yolov8_cls_trt.py
@@ -197,8 +197,10 @@ class YoLov8TRT(object):
 	    
 	    # Convert the image to row-major order, also known as "C order"
 	    image = np.ascontiguousarray(image)
+
+	    batch_data = np.expand_dims(image, axis=0)
 	    
-	    return image
+	    return batch_data
 
     def postprocess_cls(self, output_data):
         classes_ls = []