YOLO 算法的自定义数据集制作及模型训练方法

放风筝的猪

已于 2024-09-08 23:24:11 修改

阅读量508

点赞数 20

文章标签： YOLO 目标跟踪人工智能

于 2024-08-28 23:36:28 首次发布

本文链接：https://blog.csdn.net/weixin_45897172/article/details/141651199

版权

本节详细描述了如果自制数据集并用其开始训练和目标检测

1. 数据准备

这里需要用我们上一节已经标注好的数据，首先新建一个Mydataset文件夹（可自定义），如图所示放置.py文件，两个数据集文件夹以及一个数据加载配置文件，名称如图所示，.py中的代码下文将逐一给出，.yaml中的内容下文也会给出。

1.1 原始数据

原始数据集如图所示，存放在original_dataset文件夹（可自定义）下，包含了原始图片和利用labelme标注得到的txt文件。
在这里插入图片描述

1.2 待制作数据集（yolov8加载的数据集）

Mydataset文件夹下新建datasets夹（可自定义），再在datasets目录下新建train, val两个文件夹，并在其中分别创建images和labels文件夹。

train：存放训练时需要的图片和标签
val：存放测试时需要的图片和标签

下文将详细介绍每一个文件夹中存放的文件及其制作流程。

2. 按比例划分数据集

在Mydataset目录下新建一个文件split_train_val.py，代码作用是将标签文件划分为训练集和测试集
代码会按照9：1的比例将图片和标签划分并分别保存在train和val文件夹中

import os
import shutil
import random

# 设置文件夹路径
original_dataset_folder = 'Mydataset/original_dataset'
train_labels_folder = 'Mydataset/datasets/train/labels'
val_labels_folder = 'Mydataset/datasets/val/labels'
train_images_folder = 'Mydataset/datasets/train/images'
val_images_folder = 'Mydataset/datasets/val/images'

# 确保目标文件夹存在
os.makedirs(train_labels_folder, exist_ok=True)
os.makedirs(val_labels_folder, exist_ok=True)
os.makedirs(train_images_folder, exist_ok=True)
os.makedirs(val_images_folder, exist_ok=True)

# 获取所有 txt 文件
txt_files = [f for f in os.listdir(original_dataset_folder) if f.endswith('.txt')]

# 打乱文件列表
random.shuffle(txt_files)

# 计算划分点
split_index = int(len(txt_files) * 0.9)

# 划分训练集和验证集
train_txt_files = txt_files[:split_index]
val_txt_files = txt_files[split_index:]

# 处理训练集
for txt_file in train_txt_files:
    jpg_file = txt_file.replace('.txt', '.jpg')
    txt_src = os.path.join(original_dataset_folder, txt_file)
    jpg_src = os.path.join(original_dataset_folder, jpg_file)
    
    # 复制 txt 文件和对应的 jpg 文件到目标文件夹
    shutil.copy(txt_src, train_labels_folder)
    if os.path.exists(jpg_src):
        shutil.copy(jpg_src, train_images_folder)

# 处理验证集
for txt_file in val_txt_files:
    jpg_file = txt_file.replace('.txt', '.jpg')
    txt_src = os.path.join(original_dataset_folder, txt_file)
    jpg_src = os.path.join(original_dataset_folder, jpg_file)
    
    # 复制 txt 文件和对应的 jpg 文件到目标文件夹
    shutil.copy(txt_src, val_labels_folder)
    if os.path.exists(jpg_src):
        shutil.copy(jpg_src, val_images_folder)

print("数据集划分完成！")

至此YOLOv8目标检测数据集制作完成。

更多数据集格式请看官方文档

3. 数据集训练

只需在YOLOv8工程中创建数据加载配置文件并配置相关训练参数即可。
首先在github上下载代码。
在这里插入图片描述

首先将上面制作好的数据集复制到ultralytics-main文件夹中
在ultralytics文件夹中新建datasets文件夹
在datasets文件夹下新建一个Myvoc.yaml文件（可以自定义命名）

3.1 数据集yaml文件内容

在下面的代码中分别修改train和test的路径，names代表类别名称，nc代表一共多少类

# Ultralytics YOLO 🚀, AGPL-3.0 license
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC by University of Oxford
# Example usage: yolo train data=VOC.yaml
# parent
# ├── ultralytics
# └── datasets
#     └── VOC  ← downloads here (2.8 GB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../Mydataset
train: # train images (relative to 'path')  16551 images
  ./datasets/train
val: # val images (relative to 'path')  4952 images
 ./datasets/val

#ImageSets/
# Classes
names:
  0: xbl
  1: nom

nc: 2

3.2 模型yaml文件

在吗需要在ultralytics\models文件夹中新建v8文件夹，新建文件yolov8.yaml，这个文件是决定yolov8的backbone文件，其中的n,s,m,l,x可以自行选择

# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect

# Parameters
nc: 80  # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
  # [depth, width, max_channels]
  # n: [0.33, 0.25, 1024]  # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,   8.9 GFLOPs
  # s: [0.33, 0.50, 1024]  # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients,  28.8 GFLOPs
  # m: [0.67, 0.75, 768]   # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients,  79.3 GFLOPs
  # l: [1.00, 1.00, 512]   # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
  x: [1.00, 1.25, 512]   # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs

# YOLOv8.0n backbone
backbone:
  # [from, repeats, module, args]
  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
  - [-1, 3, C2f, [128, True]]
  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
  - [-1, 6, C2f, [256, True]]
  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
  - [-1, 6, C2f, [512, True]]
  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
  - [-1, 3, C2f, [1024, True]]
  - [-1, 1, SPPF, [1024, 5]]  # 9

# YOLOv8.0n head
head:
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
  - [-1, 3, C2f, [512]]  # 12

  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 12], 1, Concat, [1]]  # cat head P4
  - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 9], 1, Concat, [1]]  # cat head P5
  - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)

  - [[15, 18, 21], 1, Detect, [nc]]  # Detect(P3, P4, P5)

3.3 模型下载

选择yolov8不同的模型配置文件，可选yolov8s.pt、yolov8m.pt、yolov8x.pt等（提供了n、s、m、l、x版本，随着架构的增大，训练时间也是逐渐增大，需要根据自身设备硬件配置合理选择）。

❗❗❗注意model最好去github上提前下载好，放在如图所示的指定位置，如果没提前下载运行时会自动下载，但是下载速度可能非常非常慢。

yolov8官方权重下载地址：https://github.com/ultralytics/ultralytics?tab=readme-ov-file

3.4 train文件编写

在ultralytics-main文件夹中新建train.py文件，代码如下

epochs：训练过程中整个数据集的迭代次数。
imgsz：输入图片的尺寸大小。
batch：一次看完多少张图片才进行权重更新，梯度下降的mini-batch，提示爆显存就需要调小一点。

from ultralytics import YOLO

if __name__ == '__main__':
    # Load a model
    model = YOLO('yolov8.yaml').load('yolov8n.pt')  # build a new model from YAML
    model.train(data='Myvoc.yaml', epochs=500,imgsz=640, batch=4)    # imgsz=640,