用engine引擎文件在Jetson上面进行推理(YOLOv8)

1.pt文件,推理

在Jetson上可以通过predict.py脚本用.pt权重文件进行推理
脚本如下:

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('/best.pt') # select your model.pt path
    model.predict(source='/test',#推理图片的路径
                  imgsz=640,
                  save_txt=True,
                  project='runs/detect',
                  name='exp',
                  save=True,
                  # conf=0.2,
                  # visualize=True # visualize model features maps
                )

2.在Jetson上用engine文件进行推理

Jetson上用加载engine文件推理时也可以用脚本,但是对于Ubuntu系统来说,环境总是喜欢出现莫名奇妙的错误,因此脚本运行时,得需要加入需要的插件

LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1:/usr/lib/aarch64-linux-gnu/libGLdispatch.so.0 python3 engine_predict.py

engine_predict.py文件命令如下

import os
import warnings
import cv2
import numpy as np
import ctypes
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from ultralytics import YOLO

# 解决 TLS block 内存问题
os.environ['LD_PRELOAD'] = '/usr/lib/aarch64-linux-gnu/libgomp.so.1:/usr/lib/aarch64-linux-gnu/libGLdispatch.so.0'

warnings.filterwarnings('ignore')

# 显式加载 TensorRT 插件库
plugin_library = "/usr/lib/aarch64-linux-gnu/libnvinfer_plugin.so"
ctypes.CDLL(plugin_library)

# 初始化 TensorRT 插件
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt.init_libnvinfer_plugins(TRT_LOGGER, "")

if __name__ == '__main__':
    model = YOLO('/best.engine')  # 加载 TensorRT 模型
    model.predict(
        source='/test',#推理图片的路径
        imgsz=640,
        save_txt=True,
        project='runs/detect',
        name='exp',
        save=True,
        # conf=0.2,
        # visualize=True  # 可视化模型特征图
    )

运行结果如下所示:
用engine引擎文件在Jetson上面进行推理(YOLOv8)_第1张图片

你可能感兴趣的:(深度学习-硬件篇,YOLO,python,pycharm,深度学习,pytorch,嵌入式硬件)