Export ultralytics models via tensorrt on Jetson Orin Nano

export-tensorrt.py

# /// script
# dependencies = [
#   "bbos",
#   "torch",
#   "ultralytics",
#   "onnx",
#   "onnxruntime",
#   "onnxruntime-gpu",
#   "onnxslim",
# ]
# [tool.uv.sources]
# bbos = { path = "/home/bracketbot/BracketBotOS", editable = true }
# torch = { index = "jetson" }
# onnx = { index = "jetson" }
# onnxruntime = { index = "jetson" }
# onnxruntime-gpu = { index = "jetson" }
# onnxslim = { index = "jetson" }
# [[tool.uv.index]]
# name = "jetson"
# url = "https://pypi.jetson-ai-lab.io/jp6/cu126"
# explicit = true
# ///
"""Export YOLOv11 model to TensorRT format"""
import sys
sys.path.insert(0, '/home/bracketbot/BracketBotOS')
from bbos import Config

import torch
print(torch.cuda.is_available())
sys.path.insert(0, '/usr/lib/python3.10/dist-packages/tensorrt')
try:
    import tensorrt as trt  # noqa
    if not hasattr(trt, "__version__"):
        try:
            # Prefer package metadata (works with apt-installed dist-info)
            import importlib.metadata as md
            trt.__version__ = md.version("tensorrt")
        except Exception:
            # Fallback: grab from dist-info folder name
            import os, glob, re
            site = os.path.dirname(os.path.dirname(trt.__file__))
            m = next((re.search(r"tensorrt-([^/]+)\.dist-info", d)
                      for d in glob.glob(os.path.join(site, "tensorrt-*.dist-info"))), None)
            if m:
                trt.__version__ = m.group(1)
except Exception:
    pass
from ultralytics import YOLO

import argparse
import os

parser = argparse.ArgumentParser(description="Export YOLOv8/v11 model to TensorRT engine")
parser.add_argument("model", type=str, help="Model name (e.g., yolov8m.pt, yolov8m-worldv2.pt) or path to .pt file. NOTE: Use worldv2 models, not original world models!")
args = parser.parse_args()

# Check if it's a path or a model name
if os.path.exists(args.model):
    # Custom model path
    model_path = args.model
    base, _ = os.path.splitext(model_path)
    engine_path = base + ".engine"
else:
    # Default model name - will be downloaded by YOLO
    model_path = args.model
    # Save in daemon directory
    daemon_dir = os.path.dirname(os.path.abspath(__file__))
    model_name = os.path.splitext(os.path.basename(model_path))[0]
    base = os.path.join(daemon_dir, model_name)
    engine_path = base + ".engine"

CFG_D = Config('depth')
img_size = [CFG_D.height_D, CFG_D.width_D]

print(f"Loading model from {model_path} for image size: {img_size}")
model = YOLO(model_path)

# First export to ONNX (more reliable)
onnx_path = base + ".onnx"
print(f"Exporting to ONNX first: {onnx_path}")
model.export(format='onnx', imgsz=img_size, simplify=True, opset=17)

# Now use trtexec to convert ONNX to TensorRT (bypass ultralytics TRT bug)
print("Converting ONNX to TensorRT engine with trtexec...")
print(f"This will create {engine_path}")
import subprocess
import os
env = os.environ.copy()
cuda_paths = ':'.join([
    '/usr/local/cuda-12.6/lib64',
    '/usr/lib/aarch64-linux-gnu',
    '/usr/local/cuda/lib64',
    '/usr/local/cuda/targets/aarch64-linux/lib'
])
env['LD_LIBRARY_PATH'] = cuda_paths + ':' + env.get('LD_LIBRARY_PATH', '')
env['CUDA_VISIBLE_DEVICES'] = '0'
env['PATH'] = '/usr/local/cuda-12.6/bin:' + env.get('PATH', '')
cmd = [
    '/usr/src/tensorrt/bin/trtexec',
    f'--onnx={onnx_path}',
    f'--saveEngine={engine_path}',
    '--fp16',
    '--memPoolSize=workspace:4096M',
    '--device=0',
    '--verbose'
]
print(f"Running: {' '.join(cmd)}")
work_dir = os.path.dirname(onnx_path) if os.path.isabs(onnx_path) else os.getcwd()
try:
    result = subprocess.run(cmd, check=True, capture_output=False, env=env, cwd=work_dir)
    print(f"\nExport complete! TensorRT engine saved at {engine_path}.")
    print("The daemon will now load this engine file directly")
except subprocess.CalledProcessError as e:
    print(f"\n❌ TensorRT conversion failed!")
    print(f"If you see 'GEMM' or 'Assertion failed' errors, the model may not be compatible.")
    print(f"Try using: yolov8m-worldv2.pt or a regular yolov8m.pt/yolov11m.pt model instead.")
    raise
(venv)