YOLO：训练自己的样本数据集进行目标检测

news2024/9/22 3:47:35

作者：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
    	/train
	        image1.jpg
	        image2.jpg
	        ...
	    /val
	    	image11.jpg
	        image21.jpg
	        ...
    /labels
        /train
	        image1.txt
	        image2.txt
	        ...
	    /val
	        image11.txt
	        image21.txt
	        ...

数据集文件夹如下，

images里面的文件夹如下，train和val里面全是图片。图片格式可以是JPEG、PNG。

labels里面的文件夹如下，

train和val里面全是txt文件。

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 YOLO

if __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)