超详解|yoloV8训练自己的数据集|训练|测试|部署

1：部署源码

我的训练环境：
Ubuntu 18.04
python：3.9
torch：2.2.2+cu121
torchvision：0.17.2+cu121
GPU：Nvidia GeForce RTX 3090 * 2

这里提供中文的官方文档，方便小伙伴们阅读：Ultralytics官网
python环境，还请各位小伙伴自己安装好。

1.1：下载源码

github地址：https://github.com/ultralytics/ultralytics
可以去github上下载压缩包或者使用git命令下载源码

git clone https://github.com/ultralytics/ultralytics.git

1.2：安装运行环境

  pip install ultralytics==8.2.9 -i https://pypi.tuna.tsinghua.edu.cn/simple
  pip install hub-sdk -i https://pypi.tuna.tsinghua.edu.cn/simple
  
  # 安装torch-gpu版本
  conda install pytorch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 pytorch-cuda=12.1 -c pytorch -c nvidia

1.3：下载预训练模型

各位也可以在官网进行下载，我这边提取了下载链接，可以直接点击下载

model	下载
YOLOv8n	下载
YOLOv8s	下载
YOLOv8m	下载
YOLOv8l	下载
YOLOv8x	下载

1.4：测试

在项目根目录下，新建backbone文件夹，将下载好的预训练模型放到backbone文件夹。
使用以下代码测试，控制台输出结果。

from ultralytics import YOLO

# 加载预训练的YOLOv8n模型
model = YOLO('/backbone/yolov8s.pt')

# 在'bus.jpg'上运行推理，并附加参数
model.predict('/ultralytics/assets/bus.jpg', save=True, imgsz=640, conf=0.5)

1.5：可能遇见的错误

1：‘ has no attribute ‘FigureCanvas‘
降级matplotlib

pip install matplotlib==3.5.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

2：制作训练数据集

使用labelimg（labelimg免命令安装）进行标注的数据，训练集包含图片和xml标注文件，统一转成yolo数据格式。

2.1：将VOC格式的数据转为yolo格式

需要修改的地方，在代码中，我加上了todo，修改成自己实际的路径

# -*- coding:utf-8 -*-

"""
筛选公开数据集中，只包含目标数据集的图片和标签,自己标注的数据集也可以用。可以实现以下几点功能
1、筛选出自己感兴趣的标签，剔除不感兴趣的标签
2、按比例切分训练集、测试集、验证集
3、将xml标签文件转为yolo格式的标签
4、过滤无标签的标注文件
5、重命名转换后的文件名
"""

import os, sys, shutil
import cv2
import numpy as np

# 批量读取Annotations下的xml文件, 筛选感兴趣的标签数据
def read_xml(path_xml, input_path_img_dir, out_path_xml_dir, out_path_img_dir, class_count, pic_number,img_num, image_list):
    '''
    :param path_xml: 输入处理xml文件的绝对路径
    :return: 返回xml的label与box，其中label是一维的，并与box一一对应。
    '''
    import xml.etree.ElementTree as ET
    global classes
    with open(path_xml, 'r', encoding='UTF-8') as f:
        root = ET.parse(f).getroot()
    img_flag = False
    objects = root.findall('object')  # Get a list of all objects in this image.
    # 获取filename
    filename = str(root.find("filename").text).split(".")[0]  # 图片名称
    img_suffix = str(root.find("filename").text).split(".")[1]  # 图片后缀
    w = int(root.find("size").find('width').text)
    h = int(root.find("size").find('height').text)
    input_img_path = os.path.join(input_path_img_dir, filename + '.' + img_suffix)
    # 如果标签没有对应的图片文件，则直接返回
    if not os.path.exists(input_img_path):
        return pic_number, image_list
    print(path_xml)
    out_file = open(os.path.join(out_path_xml_dir, 'data%06d' %pic_number + '.txt'), 'w')
    # exit()
    # Parse the data for each object.
    for obj in objects:
        class_name = obj.find('name').text
        if class_name in classes:
            # 首先必须进行标签替换
            if class_name in replace_class.keys():
                class_name = replace_class[class_name]
            # 进行标签计数
            class_count[class_name] = class_count[class_name] + 1
            img_flag = True
            # Get the bounding box coordinates.
            bndbox = obj.find('bndbox')
            b1 = float(bndbox.find('xmin').text)  # bndbox.find('xmin').text
            b3 = float(bndbox.find('ymin').text)
            b2 = float(bndbox.find('xmax').text)
            b4 = float(bndbox.find('ymax').text)
            # 标注越界修正
            if b2 > w:
                b2 = w
            if b4 > h:
                b4 = h
            b = (b1, b2, b3, b4)
            bb = convert((w, h), b)
            cls_id = classes.index(class_name)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')  # 将xml数据标签，转成
        else:
            # 数据不做处理
            pass
    out_file.close()

    if img_flag:
        # 将图片分离出去
        out_img_path = os.path.join(out_path_img_dir, 'data%06d' % pic_number + '.' + img_suffix)
        # todo 这里是为了方便将图片和数据集移动到服务器的/home/taoxifa/data_collection路径之后，train.txt等文件的路径与服务器的一致（没试过相对路径，你们可以改成相对路径试一下）。
        txt_img_path = "/home/taoxifa/data_collection/images/" + 'data%06d' % pic_number + '.' + img_suffix
        shutil.copy(input_img_path, out_img_path)
        # 记录图片尺寸
        img = cv2.imdecode(np.fromfile(out_img_path, dtype=np.uint8), -1)
        size = img.shape
        if size not in class_count["size"]:
            class_count["size"].append(size)
        image_list.append(txt_img_path) # 如果是本机直接训练，不需要再把数据集移到服务器的话，将txt_img_path替换成out_img_path
        pic_number += 1
        img_num += 1
    return pic_number, image_list, img_num



def split_train_cal(txtsavepath, img_list, train_ratio, val_ratio, test_ratio):
    # 拆分数据集为训练集和测试集
    # 计算切分数据集的索引
    num_files = len(img_list)
    num_train = int(num_files * train_ratio)
    num_val = int(num_files * val_ratio)
    num_test = num_files - num_train - num_val

    # 分离训练集
    train_files = img_list[:num_train]

    # 分离验证集
    val_files = img_list[num_train:num_train + num_val]

    # 分离测试集
    test_files = img_list[num_train + num_val:]


    file_test = open(txtsavepath + '/test.txt', 'a')
    file_train = open(txtsavepath + '/train.txt', 'a')
    file_val = open(txtsavepath + '/val.txt', 'a')

    for i in train_files:
        file_train.write(i + '\n')

    for i in test_files:
        file_test.write(i + '\n')

    for i in val_files:
        file_val.write(i + '\n')

    # file_trainval.close()
    file_train.close()
    file_val.close()
    file_test.close()



def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h



def read_xml2txt():
    # todo 处理多文件夹，目前需要救护车、渣土车数据、泥浆槽罐车数据（在这个字典中，修改自己的xml数据集路径）
    dir_all = [{
        "input_path_xml_dir": r"E:\date_collection\Objects21-vehicel数据集\annotations\train",  # 救护车数据集
        "input_path_img_dir": r"E:\date_collection\Objects21-vehicel数据集\images\train"
    },{
        "input_path_xml_dir": r"E:\date_collection\渣土车\Z-渣土车数据集2400VOC\Annotations",  # 渣土车数据集
        "input_path_img_dir": r"E:\date_collection\渣土车\Z-渣土车数据集2400VOC\JPEGImages"
    },{
        "input_path_xml_dir": r"E:\date_collection\渣土车\自制渣土车数据",  # 渣土车数据集
        "input_path_img_dir": r"E:\date_collection\渣土车\自制渣土车数据"
    },{
        "input_path_xml_dir": r"E:\date_collection\槽罐车\泥浆槽罐车",  # 泥浆槽罐车数据集
        "input_path_img_dir": r"E:\date_collection\槽罐车\泥浆槽罐车"
    }]

    class_count = {"ambulance": 0, "mud_tank_truck": 0, "slag_car": 0, "size": []}
    pic_number = 7830
    for dir_path in dir_all:
        image_list = []
        input_path_xml_dir = dir_path["input_path_xml_dir"]
        input_path_img_dir = dir_path["input_path_img_dir"]
        img_num = 1
        for xml_name in os.listdir(input_path_xml_dir):
            xml_path = os.path.join(input_path_xml_dir, xml_name)
            # 后缀不是xml的跳过
            if xml_name.split('.')[1] != "xml":
                continue
            # 数量限制，先取4000张图片用于训练
            if img_num >= 4000:
                continue
            pic_number, image_list, img_num = read_xml(xml_path, input_path_img_dir, out_path_xml_dir, out_path_img_dir, class_count, pic_number,img_num, image_list)

        # 切分数据集
        split_train_cal(txtsavepath, image_list, train_ratio, val_ratio, test_ratio)

    # 将最终的处理标签结果保存下来
    print("class_count", class_count)
    # with open(os.path.join(out_path_xml_dir + "/class_count.txt"), 'w') as v:
    #     v.write(class_count)


if __name__ == "__main__":
    # 待处理的xml文件路径列表
    train_ratio = 0.7  # 训练集比例
    val_ratio = 0.2  # 验证集比例
    test_ratio = 0.1  # 测试集比例
    classes = ["ambulance", "mud_tank_truck", "slag_car", "muck-truck"]  #  todo 待提取的感兴趣的标签，不在这个列表的标签，会进行剔除
    replace_class = {"muck-truck": "slag_car"}  #  todo 替换修改标签，如果有多数据集来源，标签名称不一致，可以在这个字典中替换{"替换前的标签":"替换后的标签"}
    out_path_xml_dir = r"E:\date_collection\工业区训练数据\数据增强\labels"  #  todo 标签保存的文件夹路径
    out_path_img_dir = r"E:\date_collection\工业区训练数据\数据增强\images"  #  todo 图片保存的文件夹路径
    txtsavepath = r"E:\date_collection\工业区训练数据\数据增强"  #  todo 数据集拆分的保存路径
    if not os.path.exists(out_path_xml_dir):
        os.mkdir(out_path_xml_dir)
    if not os.path.exists(out_path_img_dir):
        os.mkdir(out_path_img_dir)
    if not os.path.exists(txtsavepath):
        os.mkdir(txtsavepath)
    
    with open(os.path.join(out_path_xml_dir + "/classes.txt"), 'w') as f:
        for i in classes:
            f.write(i + '\n')

    read_xml2txt()

执行完成之后，会在目标文件夹新增以下文件

images：存储图片
labels：存储标签
test.txt：测试数据集路径
train.txt：训练数据集路径
val.txt：验证数据集路径
！！！务必保证txt文件中的路径与实际一致，如下图所示！！！

3：修改配置文件

3.1：配置自己的数据文件

在项目路径下\ultralytics\cfg\datasets\ 新建立一个my_data.yaml文件，内容如下

path: /home/taoxifa/data_collection # dataset root dir
train: train.txt # train images (relative to 'path') 118287 images
val: val.txt # val images (relative to 'path') 5000 images
test: test.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794

# Classes
names:
  0: ambulance
  1: mud_tank_truck
  2: slag_car

3.2：修改yolov8.yaml

文件路径：\ultralytics\cfg\models\v8\yolov8.yaml，修改nc=自己的类别

4：开始训练

训练代码如下

import os
# os.environ["OMP_NUM_THREADS"]='2'

from ultralytics import YOLO

if __name__ == "__main__":
    # Load a model
    model_yaml = r"/home/taoxifa/Ai_project/yolov8/ultralytics/cfg/models/v8/yolov8s.yaml"
    data_yaml = r"/home/taoxifa/Ai_project/yolov8/ultralytics/cfg/datasets/my_data.yaml"
    pre_model = r"/home/taoxifa/Ai_project/yolov8/backbone/yolov8s.pt"

    model = YOLO(model_yaml, task='detect').load(pre_model)  # build from YAML and transfer weights
    # model = YOLO(pre_model, task='detect')  # load a pretrained model (recommended for training)

    # Train the model
    results = model.train(data=data_yaml, epochs=2000, imgsz=640, batch=16, workers=8, device=[0,1],
                          cos_lr=True, close_mosaic=200, warmup_epochs=10)

参数解释
data：数据集配置文件的路径，指定用于训练的数据集配置文件
epochs：训练过程中整个数据集将被迭代多少次。
imgsz：用于设置输入图像尺寸。训练时，只能是整数32的倍数，测试和预测时，可以使用设置[1920,1080]。
batch：每个批次中的图像数量。
workers：用于设置数据加载过程中的线程数
device：模型运行的设置，多gpu时，用列表表示
cos_lr：余弦学习率调度器，设置为True可以帮助模型在训练过程中按照余弦函数的形状调整学习率，从而在训练初期使用较高的学习率，有助于快速收敛，而在训练后期逐渐降低学习率，有助于细致调整模型参数
close_mosaic：用于确定是否在最后几个训练周期中禁用马赛克数据增强
warmup_epochs：预热学习轮数，学习率从低值逐渐增加到初始学习率，以在早期稳定训练
更多参数，参考：ultralytics-训练参数

5：查看训练结果

训练结果存储在runs/detect/train下，可以查看评估的结果和最优模型。
模型在weights路径下的best.pt。下面的图为训练15轮后的结果，非2000轮后的结果。

6：模型评估

代码评估

from ultralytics import YOLO

if __name__ == '__main__':

    # Load a model
    model_yaml = r"/home/taoxifa/Ai_project/yolov8/ultralytics/cfg/models/v8/yolov8s.yaml"
    data_yaml = r"/home/taoxifa/Ai_project/yolov8/ultralytics/cfg/datasets/my_data.yaml"
    pre_model = r"/home/taoxifa/Ai_project/yolov8/runs/detect/train/weights/best.pt"

    # model = YOLO(model_yaml, task='detect').load(pre_model)  # build from YAML and transfer weights
    model = YOLO(pre_model, task='detect')  # load a pretrained model (recommended for training)

    # Train the model
    results = model.val(data=data_yaml, imgsz=640, device=[0, 1])

指令评估

yolo detect val model=/home/taoxifa/Ai_project/yolov8/runs/detect/train/weights/best.pt data=/home/taoxifa/Ai_project/yolov8/ultralytics/cfg/datasets/my_data.yaml

评估结果

7：模型预测

from ultralytics import YOLO



if __name__ == '__main__':

    pth_path = r"/home/taoxifa/Ai_project/yolov8/runs/detect/train/weights/best.pt"

    test_path = r"/home/taoxifa/Ai_project/yolov8/images"
    # Load a model
    # model = YOLO('yolov8n.pt')  # load an official model
    model = YOLO(pth_path)  # load a custom model

    # Predict with the model
    results = model(test_path, save=True, conf=0.6, iou=0.8)  # predict on an image

输出图片

8：模型导出（导出为onnx，可用于tensorrt或openvino推理部署使用）

安装其他依赖环境

  pip install openvino-dev==2024.0.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
  pip install onnx==1.16.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
  pip install nncf==2.9.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
  pip install onnxruntime==1.17.3 -i https://pypi.tuna.tsinghua.edu.cn/simple

from ultralytics import YOLO

if __name__ == "__main__":

    pth_path = r"/home/taoxifa/Ai_project/yolov8/runs/detect/train/weights/best.pt"

    # Load a model
    model = YOLO("/home/taoxifa/Ai_project/yolov8/backbone/yolov8s.pt")  # 加载官方模型
    model = YOLO(pth_path)  # 加载自定义训练模型

    # Export the model
    model.export(format='openvino', int8=True) # 设置int8，会自动下载coco128的数据集对齐，量化成int8模型

9：openvino Or ONNX 推理

openvino官网：官网openvino教程

# 官方位置：https://docs.openvino.ai/2024/notebooks/004-hello-detection-with-output.html

import argparse
import time
import cv2
import numpy as np
from openvino.runtime import Core  # pip install openvino -i  https://pypi.tuna.tsinghua.edu.cn/simple
import onnxruntime as ort  # 使用onnxruntime推理用上，pip install onnxruntime，默认安装CPU

# COCO默认的80类
CLASSES = ['ambulance', 'mud_tank_truck', 'slag_car']


class OpenvinoInference(object):
    def __init__(self, onnx_path):
        self.onnx_path = onnx_path
        ie = Core()
        self.model_onnx = ie.read_model(model=self.onnx_path)
        self.compiled_model_onnx = ie.compile_model(model=self.model_onnx, device_name="CPU", config={})
        self.output_layer_onnx = self.compiled_model_onnx.output(0)

    def predict(self, datas):
        predict_data = self.compiled_model_onnx([datas])[self.output_layer_onnx]
        return predict_data


class YOLOv8:
    """YOLOv8 object detection model class for handling inference and visualization."""

    def __init__(self, onnx_model, imgsz=(640, 640), infer_tool='openvino'):
        """
        Initialization.

        Args:
            onnx_model (str): Path to the ONNX model.
        """
        self.infer_tool = infer_tool
        if self.infer_tool == 'openvino':
            # 构建openvino推理引擎
            self.openvino = OpenvinoInference(onnx_model)
            self.ndtype = np.single
        else:
            # 构建onnxruntime推理引擎
            self.ort_session = ort.InferenceSession(onnx_model,
                                                    providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
                                                    if ort.get_device() == 'GPU' else ['CPUExecutionProvider'])

            # Numpy dtype: support both FP32 and FP16 onnx model
            self.ndtype = np.half if self.ort_session.get_inputs()[0].type == 'tensor(float16)' else np.single

        self.classes = CLASSES  # 加载模型类别
        self.model_height, self.model_width = imgsz[0], imgsz[1]  # 图像resize大小
        self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))  # 为每个类别生成调色板

    def __call__(self, im0, conf_threshold=0.4, iou_threshold=0.45):
        """
        The whole pipeline: pre-process -> inference -> post-process.

        Args:
            im0 (Numpy.ndarray): original input image.
            conf_threshold (float): confidence threshold for filtering predictions.
            iou_threshold (float): iou threshold for NMS.

        Returns:
            boxes (List): list of bounding boxes.
        """
        # 前处理Pre-process
        t1 = time.time()
        im, ratio, (pad_w, pad_h) = self.preprocess(im0)
        print('预处理时间：{:.3f}s'.format(time.time() - t1))

        # 推理 inference
        t2 = time.time()
        if self.infer_tool == 'openvino':
            preds = self.openvino.predict(im)
        else:
            preds = self.ort_session.run(None, {self.ort_session.get_inputs()[0].name: im})[0]
        print('推理时间：{:.2f}s'.format(time.time() - t2))

        # 后处理Post-process
        t3 = time.time()
        boxes = self.postprocess(preds,
                                 im0=im0,
                                 ratio=ratio,
                                 pad_w=pad_w,
                                 pad_h=pad_h,
                                 conf_threshold=conf_threshold,
                                 iou_threshold=iou_threshold,
                                 )
        print('后处理时间：{:.3f}s'.format(time.time() - t3))

        return boxes

    # 前处理，包括：resize, pad, HWC to CHW，BGR to RGB，归一化，增加维度CHW -> BCHW
    def preprocess(self, img):
        """
        Pre-processes the input image.

        Args:
            img (Numpy.ndarray): image about to be processed.

        Returns:
            img_process (Numpy.ndarray): image preprocessed for inference.
            ratio (tuple): width, height ratios in letterbox.
            pad_w (float): width padding in letterbox.
            pad_h (float): height padding in letterbox.
        """
        # Resize and pad input image using letterbox() (Borrowed from Ultralytics)
        shape = img.shape[:2]  # original image shape
        new_shape = (self.model_height, self.model_width)
        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        ratio = r, r
        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
        pad_w, pad_h = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2  # wh padding
        if shape[::-1] != new_unpad:  # resize
            img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
        top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
        left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
        img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))  # 填充
        # Transforms: HWC to CHW -> BGR to RGB -> div(255) -> contiguous -> add axis(optional)
        img = np.ascontiguousarray(np.einsum('HWC->CHW', img)[::-1], dtype=self.ndtype) / 255.0
        img_process = img[None] if len(img.shape) == 3 else img
        return img_process, ratio, (pad_w, pad_h)

    # 后处理，包括：阈值过滤与NMS
    def postprocess(self, preds, im0, ratio, pad_w, pad_h, conf_threshold, iou_threshold):
        """
        Post-process the prediction.

        Args:
            preds (Numpy.ndarray): predictions come from ort.session.run().
            im0 (Numpy.ndarray): [h, w, c] original input image.
            ratio (tuple): width, height ratios in letterbox.
            pad_w (float): width padding in letterbox.
            pad_h (float): height padding in letterbox.
            conf_threshold (float): conf threshold.
            iou_threshold (float): iou threshold.

        Returns:
            boxes (List): list of bounding boxes.
        """
        x = preds  # outputs: predictions (1, 84, 8400)
        # Transpose the first output: (Batch_size, xywh_conf_cls, Num_anchors) -> (Batch_size, Num_anchors, xywh_conf_cls)
        x = np.einsum('bcn->bnc', x)  # (1, 8400, 84)

        # Predictions filtering by conf-threshold
        x = x[np.amax(x[..., 4:], axis=-1) > conf_threshold]

        # Create a new matrix which merge these(box, score, cls) into one
        # For more details about `numpy.c_()`: https://numpy.org/doc/1.26/reference/generated/numpy.c_.html
        x = np.c_[x[..., :4], np.amax(x[..., 4:], axis=-1), np.argmax(x[..., 4:], axis=-1)]

        # NMS filtering
        # 经过NMS后的值, np.array([[x, y, w, h, conf, cls], ...]), shape=(-1, 4 + 1 + 1)
        x = x[cv2.dnn.NMSBoxes(x[:, :4], x[:, 4], conf_threshold, iou_threshold)]

        # 重新缩放边界框，为画图做准备
        if len(x) > 0:
            # Bounding boxes format change: cxcywh -> xyxy
            x[..., [0, 1]] -= x[..., [2, 3]] / 2
            x[..., [2, 3]] += x[..., [0, 1]]

            # Rescales bounding boxes from model shape(model_height, model_width) to the shape of original image
            x[..., :4] -= [pad_w, pad_h, pad_w, pad_h]
            x[..., :4] /= min(ratio)

            # Bounding boxes boundary clamp
            x[..., [0, 2]] = x[:, [0, 2]].clip(0, im0.shape[1])
            x[..., [1, 3]] = x[:, [1, 3]].clip(0, im0.shape[0])

            return x[..., :6]  # boxes
        else:
            return []


    # 绘框
    def draw_and_visualize(self, im, bboxes, vis=False, save=True):
        """
        Draw and visualize results.

        Args:
            im (np.ndarray): original image, shape [h, w, c].
            bboxes (numpy.ndarray): [n, 6], n is number of bboxes.
            vis (bool): imshow using OpenCV.
            save (bool): save image annotated.

        Returns:
            None
        """
        # Draw rectangles
        for (*box, conf, cls_) in bboxes:
            # draw bbox rectangle
            cv2.rectangle(im, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])),
                          self.color_palette[int(cls_)], 1, cv2.LINE_AA)
            cv2.putText(im, f'{self.classes[int(cls_)]}: {conf:.3f}', (int(box[0]), int(box[1] - 9)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, self.color_palette[int(cls_)], 2, cv2.LINE_AA)

        # # Show image
        # if vis:
        #     cv2.imshow('demo', im)
        #     cv2.waitKey(1)


        # # Save image
        if save:
            print("保存图片成功")
            cv2.imwrite('demo.png', im)
        return im

def read_Video(video_path, model, conf_threshold, iou_threshold):
    cap = cv2.VideoCapture(video_path)


    while True:
        # Capture NEXT frame
        # next_frame = player.next()
        start_time = time.time()
        frame_number = 0
        ret, img = cap.read()
        if ret:
            boxes = model(img, conf_threshold, iou_threshold)
            if len(boxes) > 0:
                img = model.draw_and_visualize(img, boxes, vis=False, save=True)
            stop_time = time.time()
            total_time = stop_time - start_time
            frame_number = frame_number + 1
            async_fps = frame_number / total_time
            print("async_fps",async_fps)
            cv2.putText(img, f'{async_fps:.3f}', (20, 20),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2, cv2.LINE_AA)
            cv2.imshow("img", img)
            key = cv2.waitKey(1)
            # escape = 27
            if key == 27:
                break


if __name__ == '__main__':
    # Create an argument parser to handle command-line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default=r'/home/taoxifa/Ai_project/yolov8/runs/detect/train/weights/best_int8_openvino_model/best.xml', help='Path to ONNX model')
    parser.add_argument('--source', type=str, default=str(r'/home/taoxifa/Ai_project/yolov8/images/data007753.jpg'), help='Path to input image')
    parser.add_argument('--imgsz', type=tuple, default=(640, 640), help='Image input size')
    parser.add_argument('--conf', type=float, default=0.6, help='Confidence threshold')
    parser.add_argument('--iou', type=float, default=0.1, help='NMS IoU threshold')
    parser.add_argument('--infer_tool', type=str, default='openvino', choices=("openvino", "onnxruntime"),
                        help='选择推理引擎')
    args = parser.parse_args()

    # Build model
    model = YOLOv8(args.model, args.imgsz, args.infer_tool)

    # Read image by OpenCV
    img = cv2.imread(args.source)
    # 视频推理
    # video_path = r"E:\date_collection\安全帽\工地全景视频.mp4"
    # read_Video(video_path, model, conf_threshold=args.conf, iou_threshold=args.iou)  # 视频推理


    # Inference
    boxes = model(img, conf_threshold=args.conf, iou_threshold=args.iou) # 图片推理

    # Visualize
    if len(boxes) > 0:
        model.draw_and_visualize(img, boxes, vis=False, save=True)
    # cv2.destroyAllWindows()

码字不易，劳烦点赞+收藏。有问题，可以在评论区讨论，都会一一回应，谢谢。