haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update yolov8_cls_trt.py (#1455)

Browse Source 
batch image data
This commit is contained in:
lindsayshuo 2024-03-14 13:53:42 +08:00 committed by GitHub
parent 5d7896913c
commit c1b5cfdf2d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
yolov8/yolov8_cls_trt.py
View File

@ -197,8 +197,10 @@ class YoLov8TRT(object):
# Convert the image to row-major order, also known as "C order" # Convert the image to row-major order, also known as "C order"
image = np.ascontiguousarray(image) image = np.ascontiguousarray(image)
batch_data = np.expand_dims(image, axis=0)
return image return batch_data
def postprocess_cls(self, output_data): def postprocess_cls(self, output_data):
classes_ls = [] classes_ls = []