YOLOv8训练自己的数据集并推理

 训练

1、 下载yolov8源码

yolo官网官网源码下载地址： yolo官网

如果打不开的话可以科学上网。

2、处理数据集（转换与划分）

数据集VOC格式转yolo格式
如何查看自己数据集格式，打开Annotations文件夹，如果看到文件后缀为.xml，则为VOC格式，如果文件后缀为.txt则为yolo格式。yolov8训练需要转为yolo格式训练，转换代码如下。

import os
import xml.etree.ElementTree as ET

# 定义类别顺序
categories = ['name1', 'name2']
category_to_index = {category: index for index, category in enumerate(categories)}

# 定义输入文件夹和输出文件夹
input_folder = r'D:\临时存放\数据集\VOC2028\Annotations'  # 替换为实际的XML文件夹路径
output_folder = r'D:\临时存放\数据集\VOC2028\label'  # 替换为实际的输出TXT文件夹路径

# 确保输出文件夹存在
os.makedirs(output_folder, exist_ok=True)

# 遍历输入文件夹中的所有XML文件
for filename in os.listdir(input_folder):
    if filename.endswith('.xml'):
        xml_path = os.path.join(input_folder, filename)
        # 解析XML文件
        tree = ET.parse(xml_path)
        root = tree.getroot()
        # 提取图像的尺寸
        size = root.find('size')
        width = int(size.find('width').text)
        height = int(size.find('height').text)
        # 存储name和对应的归一化坐标
        objects = []

        # 遍历XML中的object标签
        for obj in root.findall('object'):
            name = obj.find('name').text
            if name in category_to_index:
                category_index = category_to_index[name]
            else:
                continue  # 如果name不在指定类别中，跳过该object

            bndbox = obj.find('bndbox')
            xmin = int(bndbox.find('xmin').text)
            ymin = int(bndbox.find('ymin').text)
            xmax = int(bndbox.find('xmax').text)
            ymax = int(bndbox.find('ymax').text)

            # 转换为中心点坐标和宽高
            x_center = (xmin + xmax) / 2.0
            y_center = (ymin + ymax) / 2.0
            w = xmax - xmin
            h = ymax - ymin

            # 归一化
            x = x_center / width
            y = y_center / height
            w = w / width
            h = h / height

            objects.append(f"{category_index} {x} {y} {w} {h}")

        # 输出结果到对应的TXT文件
        txt_filename = os.path.splitext(filename)[0] + '.txt'
        txt_path = os.path.join(output_folder, txt_filename)
        with open(txt_path, 'w') as f:
            for obj in objects:
                f.write(obj + '\n')

数据集切分成 训练集、测试集和验证集

我代码中切分的比例是7：2：1，可以根据自己需求划分。代码如下。

import os
import shutil
import random


# random.seed(0)  #随机种子，可自选开启
def split_data(file_path, label_path, new_file_path, train_rate, val_rate, test_rate):
    images = os.listdir(file_path)
    labels = os.listdir(label_path)
    images_no_ext = {os.path.splitext(image)[0]: image for image in images}
    labels_no_ext = {os.path.splitext(label)[0]: label for label in labels}
    matched_data = [(img, images_no_ext[img], labels_no_ext[img]) for img in images_no_ext if img in labels_no_ext]

    unmatched_images = [img for img in images_no_ext if img not in labels_no_ext]
    unmatched_labels = [label for label in labels_no_ext if label not in images_no_ext]
    if unmatched_images:
        print("未匹配的图片文件:")
        for img in unmatched_images:
            print(images_no_ext[img])
    if unmatched_labels:
        print("未匹配的标签文件:")
        for label in unmatched_labels:
            print(labels_no_ext[label])

    random.shuffle(matched_data)
    total = len(matched_data)
    train_data = matched_data[:int(train_rate * total)]
    val_data = matched_data[int(train_rate * total):int((train_rate + val_rate) * total)]
    test_data = matched_data[int((train_rate + val_rate) * total):]

    # 处理训练集
    for img_name, img_file, label_file in train_data:
        old_img_path = os.path.join(file_path, img_file)
        old_label_path = os.path.join(label_path, label_file)
        new_img_dir = os.path.join(new_file_path, 'train', 'images')
        new_label_dir = os.path.join(new_file_path, 'train', 'labels')
        os.makedirs(new_img_dir, exist_ok=True)
        os.makedirs(new_label_dir, exist_ok=True)
        shutil.copy(old_img_path, os.path.join(new_img_dir, img_file))
        shutil.copy(old_label_path, os.path.join(new_label_dir, label_file))
    # 处理验证集
    for img_name, img_file, label_file in val_data:
        old_img_path = os.path.join(file_path, img_file)
        old_label_path = os.path.join(label_path, label_file)
        new_img_dir = os.path.join(new_file_path, 'val', 'images')
        new_label_dir = os.path.join(new_file_path, 'val', 'labels')
        os.makedirs(new_img_dir, exist_ok=True)
        os.makedirs(new_label_dir, exist_ok=True)
        shutil.copy(old_img_path, os.path.join(new_img_dir, img_file))
        shutil.copy(old_label_path, os.path.join(new_label_dir, label_file))
    # 处理测试集
    for img_name, img_file, label_file in test_data:
        old_img_path = os.path.join(file_path, img_file)
        old_label_path = os.path.join(label_path, label_file)
        new_img_dir = os.path.join(new_file_path, 'test', 'images')
        new_label_dir = os.path.join(new_file_path, 'test', 'labels')
        os.makedirs(new_img_dir, exist_ok=True)
        os.makedirs(new_label_dir, exist_ok=True)
        shutil.copy(old_img_path, os.path.join(new_img_dir, img_file))
        shutil.copy(old_label_path, os.path.join(new_label_dir, label_file))
    print("数据集已划分完成")


if __name__ == '__main__':
    file_path = r"D:\临时存放\数据集\VOC2028\JPEGImages"  # 图片文件夹
    label_path = r'D:\临时存放\数据集\VOC2028\label'  # 标签文件夹
    new_file_path = r"D:\临时存放\数据集\helmet_det"  # 新数据存放位置
    split_data(file_path, label_path, new_file_path, train_rate=0.7, val_rate=0.2, test_rate=0.1)

3、创建data.yaml

在下载的源码根目录下创建一个新的data.yaml文件内容如下。主要是数据集路径和标签信息，根据自己的数据集去替换。

train: D:/临时存放/数据集/helmet_det/train/images  
val: D:/临时存放/数据集/helmet_det/val/images  
test: D:/临时存放/数据集/helmet_det/test/images

nc: 2 #你的数据集标签类别数

# Classes
names: ['name1','name2']

4、创建一个自己的yolo8.yaml

打开yolov8源码，找到路径 ultralytics-main\ultralytics\cfg\models\8\yolo8.yaml。复制该文件内容，创建一个yolo8_copy.yaml文件，也可命名其他的名字后缀为.yaml就行，内容如下。

# 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)

5、下载预训练模型

 yolo官网地址： yolo官网

模型大小n<s<m<l<x，需要用到那个型号的模型直接下载就行。

6、开始训练

训练代码如下，epochs是训练轮数，可以由少变多看训练效果，workers和batch根据电脑性能进行调整。

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

import warnings
from re import search

warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':

    model = YOLO(model=r'C:\Users\Desktop\yolov8\ultralytics-main\ultralytics\cfg\models\8\yolo8_copy.yaml')
    model.load('yolov8s.pt') # 加载预训练权重,改进或者做对比实验时候不建议打开，因为用预训练模型整体精度没有很明显的提升
    model.train(data=r'data.yaml',
                imgsz=640,
                epochs=120,
                batch=16,
                workers=0,
                device='',
                optimizer='SGD',
                close_mosaic=10,
                resume=False,
                project='runs/train',
                name='exp',
                single_cls=False,
                cache=False,
                )

推理

推理代码如下，记得替换自己的模型和图片路径。

from ultralytics import YOLO
# 加载训练好的模型，改为自己的路径
model = YOLO('runs/train/exp/weights/best.pt')
source = 'test.jpg' #修改为自己的图片路径及文件名
# 运行推理，并附加参数
model.predict(source, save=True)