YOLOV8权重onnx文件python部署篇

前言

尝试了多种onnx的部署代码，不是nms标的框太大就是没有类别标注，后面才发现了yolov8官方自带了部署例子，只需要更改部分代码即可！写此篇文章谨记走过的大坑。

导出onnx权重

# 导出 onnx
from ultralytics import YOLO

model = YOLO('/root/runs/detect/train3/weights/best.pt')
success = model.export(format="onnx")

这里只需要将pt文件换成自己的权重代码即可，最后导出的onnx的文件也会存到这个文件夹里

修改代码

找到官方给的案例，这个是python的部署方式，直接调用的是onnx文件

ultralytics-main/examples/YOLOv8-OpenCV-ONNX-Python/mian.py

1、第一处

import argparse

import cv2.dnn
import numpy as np

from ultralytics.utils import ASSETS, yaml_load
from ultralytics.utils.checks import check_yaml

CLASSES = yaml_load(check_yaml('/root/data.yaml'))['names']

colors = np.random.uniform(0, 255, size=(len(CLASSES), 3))

 将第9行的CLASSES换成自己的yaml

2、第二处

如果是linux系统的话，如果不改这里可能会报错，因为没有界面。如果想要保存的话就再加一行代码。

注释的代码为第68行

cv2.imshow('image', original_image)

 加的存储代码为我自己加的，仅供参考

cv2.imwrite('output_image.jpg', original_image)

3、第三处

改变参数model和img，改成自己就可以了！直接运行就可以的到检测过的图像了。

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', default='/root/ultralytics-main/weights/best.onnx', help='Input your onnx model.')
    parser.add_argument('--img', default=str(ASSETS / '/root/ultralytics-main/1.jpg'), help='Path to input image.')
    args = parser.parse_args()
    main(args.model, args.img)

结语

仅供学习参考使用！！！