Update yolov8_cls_trt.py (#1455)

batch image data
2024-03-14 13:53:42 +08:00 · 2024-03-14 13:53:42 +08:00 · c1b5cfdf2d
commit c1b5cfdf2d
parent 5d7896913c
1 changed files with 3 additions and 1 deletions
--- 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 = []