Last active
February 26, 2024 17:24
-
-
Save adujardin/5f2b2a1474ebb1814e068ba84fd8e8fb to your computer and use it in GitHub Desktop.
YOLOv8 TensorRT Python
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# From https://docs.ultralytics.com/fr/integrations/tensorrt/#installation | |
from ultralytics import YOLO | |
import os | |
import cv2 | |
# Export the model to TensorRT format, this will take a long time but should be done only once (per GPU) | |
if not os.path.isfile('yolov8n.engine'): | |
# Load the YOLOv8 model | |
model = YOLO('yolov8n.pt') | |
# Optimize the model | |
model.export(format='engine', half=True) # creates 'yolov8n.engine' | |
# Load the exported TensorRT model | |
tensorrt_model = YOLO('yolov8n.engine', task="detect") | |
# Run inference | |
# Directly from downloaded image | |
results = tensorrt_model('https://ultralytics.com/images/bus.jpg') | |
# From OpenCV Mat | |
im = cv2.imread("bus.jpg") | |
results = tensorrt_model(im) | |
print(results) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment