作者:CSDN @ _养乐多_
本文将介绍如何使用python语言和 ultralytics 库训练自己的数据集,并进行 YOLO 目标检测模型训练和推理的代码。
文章目录
- 一、样本数据集准备
- 1.1 标注工具
- 1.2 数据集格式
- 1.2.1 图片和标签数据集制作
- 1.2.2 data.yaml制作
- 二、模型训练
- 三、图片推理
一、样本数据集准备
1.1 标注工具
labelimg 、labelme 、VIA 都可以
1.2 数据集格式
本文直接采用YOLO数据集格式进行示例,
1.2.1 图片和标签数据集制作
YOLO数据集文件结构如下:
/dataset/images/trainimage1.jpgimage2.jpg.../valimage11.jpgimage21.jpg.../labels/trainimage1.txtimage2.txt.../valimage11.txtimage21.txt...
数据集文件夹如下,
images里面的文件夹如下,train和val里面全是图片。图片格式可以是JPEG、PNG。
txt中的内容如下图所示,格式为class-id x y width height
,代表
- class_id: 类别ID,从0开始,自己定义
- center_x:边界框中心的x坐标,相对于图像宽度进行归一化(值在0到1之间)
- center_y:边界框中心的y坐标,相对于图像高度进行归一化(值在0到1之间)
- width:边界框的宽度,相对于图像宽度进行归一化(值在0到1之间)
- height:边界框的高度,相对于图像高度进行归一化(值在0到1之间)
1.2.2 data.yaml制作
写一个data.yaml文件,并保存。
train: E:\\\\images\\train # 训练数据集路径
val: E:\\images\\val # 验证数据集路径
nc: 1 # 类别数量
names: ['bus'] # 类别名字
注意:labels为啥没有路径,因为它可以自己识别,只要和images同级目录就行。
二、模型训练
from ultralytics import YOLOif __name__ == '__main__':# 加载模型model = YOLO("yolov8n.pt") # 加载一个预训练模型(推荐使用预训练模型用于训练)# 训练model.train(data="data.yaml", epochs=3, workers=0) # 使用单进程数据加载训练模型# 验证metrics = model.val() # 评估模型在验证集上的性能# 保存模型path = model.export(format="onnx") # 将模型导出为ONNX格式
三、图片推理
from ultralytics import YOLO# 加载训练好的模型
model = YOLO("./runs/detect/train40/weights/best.pt")# 读取图片
image_path = "./images/train/0001.jpg"model.predict(image_path, imgsz=640, save=True, device=0)
测试图片,
=
=
=
=
=
=
=
=
=
=
=
=
=
=
=
=
=
=