YOLO：目标检测模型的训练和推理（简单Demo）

YOLO：目标检测模型的训练和推理（简单Demo）

引用

CSDN

1.

https://blog.csdn.net/qq_35591253/article/details/140846683

本文将通过一个简单的Demo，详细介绍如何使用YOLO模型进行目标检测，包括模型训练、推理以及环境配置等内容。

一、准备

1.1 模型类型

YOLO8模型性能数据：

这个表格包含了YOLO模型的各项性能指标，包括模型大小、mAP值、CPU和TensorRT的推理速度、参数数量以及FLOPs。

1.2 环境配置

安装ultralytics库：

pip install ultralytics

二、模型训练和推理

Windows系统下，直接运行以下代码。保持网络通畅，因为模型文件和推理图片都是通过网络直接下载并应用的。

from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型
    # model = YOLO("yolov8n.yaml")  # 从头构建一个新模型
    model = YOLO("yolov8n.pt")  # 加载一个预训练模型（推荐用于训练）
    # 训练
    model.train(data="coco8.yaml", epochs=3, workers=0)  # 使用单进程数据训练模型
    # 验证
    metrics = model.val()  # 评估模型在验证集上的性能
    
    # 推理
    # results = model("https://ultralytics.com/images/bus.jpg")  # 对图像进行预测
    model.predict("https://ultralytics.com/images/bus.jpg", imgsz=640, save=True, device=0)
    # 保存模型
    path = model.export(format="onnx")  # 将模型导出为ONNX格式

如果网络有问题，就点击链接下载bus.jpg图片。也可以自己在网上找一个有很多物体的图片进行推理。

代码如下：

from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型
    # model = YOLO("yolov8n.yaml")  # 从头构建一个新模型
    model = YOLO("yolov8n.pt")  # 加载一个预训练模型（推荐用于训练）
    # 训练
    model.train(data="coco8.yaml", epochs=3, workers=0)  # 使用单进程数据训练模型
    # 验证
    metrics = model.val()  # 评估模型在验证集上的性能
    
    # 推理
    # 如果网络有问题，就点击链接下载到本地，“https://ultralytics.com/images/bus.jpg”
    # results = model("bus.jpg")  # 对图像进行预测
    model.predict("bus.jpg", imgsz=640, save=True, device=0)
    # 保存模型
    path = model.export(format="onnx")  # 将模型导出为ONNX格式

三、推理

3.1 基于pt格式模型

修改训练好的模型路径，图片的路径。只需4行，即可实现推理。

from ultralytics import YOLO

model = YOLO("../模型/best.pt")
image_path = "../测试图片/2.png"
model.predict(image_path, imgsz=640, save=True, device=0)

3.2 基于onnx格式模型

from ultralytics import YOLO

# 加载你的模型
# model = YOLO("yolov8n.pt")
# # 导出模型
# model.export(format="onnx")  # creates 'yolov8n.onnx'
# 模型推理
onnx_model = YOLO("./model/yolov8n.onnx")
# 选择你的图片、视频
results = onnx_model("./images/0_frame_001.jpg")
result = results[0]
result.show()  # 显示推理结果

四、讨论

由于只训练了3次，训练的epoch较少，精度就会低一点。例如，在下面的图片中，右上角的人和马都检测错了。

再简化一下代码，四行就可以训练、预测和保存模型。

from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型
    model = YOLO("yolov8n.pt")  # 加载一个预训练模型（推荐用于训练）
    # 训练
    model.train(data="coco8.yaml", epochs=3, workers=0)  # 使用单进程数据训练模型
    
    # 推理
    model.predict("https://ultralytics.com/images/bus.jpg", imgsz=640, save=True, device=0)
    # 保存模型
    path = model.export(format="onnx")  # 将模型导出为ONNX格式