YOLOv5与YOLOv8 训练准备工作（不包含环境搭建）

前言：我发现除了安装环境需要耗费大量时间以外，对于训练前的准备工作也要琢磨一段时间，所以本篇主要讲一下训练前需要准备的工作（主要是XML格式换为txt，以及划分数据集验证集，和训练参数的设置），我在这里都总结一下，留给下次训练的我查看，以及广大网友参考。（默认一些基础的环境例如（CUDA.pytorch.torchvision.opencv）已经搭建好以及安装完requirements）2024.8.6

YOLOv5：

XML格式换为txt，以及划分数据集验证集：

1.首先需要准备好图片(需要训练的图片)和标签(通过labelimg打完标签)按照以下目录存放：
                                         
格式如下：(需要自己完整地创造出以下目录，方便后续进行)
|——VOCdevkit
||——images
        ——train  (空)     
        ——val    (空)
||——labels
        ——train  (空)
        ——val    (空)
||——VOC2007
        ——Annotations  (放置自己的标签XML文件)
        ——JPEGImages   (放置自己的训练图片)
        ——label  (空)  

2.然后在yolov5-5.0目录下（位置不重要）创建xml转YOLO.py 复制粘贴代码↓

代码如下： （没有安装包就pip Install一下）

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join


def convert(size, box):
    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]
    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]
    return (x, y, w, h)


def convert_annotation(xml_files_path, save_txt_files_path, classes):
    xml_files = os.listdir(xml_files_path)
    print(xml_files)
    for xml_name in xml_files:
        print(xml_name)
        xml_file = os.path.join(xml_files_path, xml_name)
        out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')
        out_txt_f = open(out_txt_path, 'w')
        tree = ET.parse(xml_file)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)

        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes or int(difficult) == 1:
                continue
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            # b=(xmin, xmax, ymin, ymax)
            print(w, h, b)
            bb = convert((w, h), b)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


if __name__ == "__main__":
    # 需要转换的类别，需要一一对应  替换为自己的类别
    classes1 = ['red', 'green','blue','bottle']  
    # 2、voc格式的xml标签文件路径 替换为自己的标签文件文件xml（Annotations）所在绝对路径
    xml_files1 = r'D:\yolov5\yolov5-5.0\yolov5-5.0\VOCdevkit\VOC2007\Annotations'
    # 3、转化为yolo格式的txt标签文件存储路径  替换为自己想要储存txt标签文件的路径 （label） 
    save_txt_files1 = r'D:\yolov5\yolov5-5.0\yolov5-5.0\VOCdevkit\VOC2007\label'

    convert_annotation(xml_files1, save_txt_files1, classes1)

注：修改一下classes1    xml_files1    save_txt_files1 替换为自己的路径

运行过后会得到自己的txt文件：

                                                      

3.然后把JPEGImages 里的图片和刚刚转换好的label里的文件分别复制到 images/train || labels/train这两个地方：

                            

4.然后在yolov5-5.0目录下（位置不重要）创建划分数据集和验证集.py  复制粘贴代码↓

import os, random, shutil


def moveimg(fileDir, tarDir):
    pathDir = os.listdir(fileDir)  # 取图片的原始路径
    filenumber = len(pathDir)
    rate = 0.1  # 自定义抽取图片的比例，比方说100张抽10张，那就是0.1
    picknumber = int(filenumber * rate)  # 按照rate比例从文件夹中取一定数量图片
    sample = random.sample(pathDir, picknumber)  # 随机选取picknumber数量的样本图片
    print(sample)
    for name in sample:
        shutil.move(fileDir + name, tarDir + "\\" + name)
    return


def movelabel(file_list, file_label_train, file_label_val):
    for i in file_list:
        if i.endswith('.jpg'):
            # filename = file_label_train + "\\" + i[:-4] + '.xml'  # 可以改成xml文件将’.txt‘改成'.xml'就可以了
            filename = file_label_train + "\\" + i[:-4] + '.txt'  # 可以改成xml文件将’.txt‘改成'.xml'就可以了
            if os.path.exists(filename):
                shutil.move(filename, file_label_val)
                print(i + "处理成功！")


if __name__ == '__main__':
    fileDir = r"D:\yolov5\yolov5-5.0\yolov5-5.0\VOCdevkit\images\train" + "\\"  # 源图片文件夹路径
    tarDir = r'D:\yolov5\yolov5-5.0\yolov5-5.0\VOCdevkit\images\val'  # 图片移动到新的文件夹路径
    moveimg(fileDir, tarDir)
    file_list = os.listdir(tarDir)
    file_label_train = r"D:\yolov5\yolov5-5.0\yolov5-5.0\VOCdevkit\labels\train"  # 源图片标签路径
    file_label_val = r"D:\yolov5\yolov5-5.0\yolov5-5.0\VOCdevkit\labels\val"  # 标签
    # 移动到新的文件路径
    movelabel(file_list, file_label_train, file_label_val)

注：修改一下

fileDir   = ——\VOCdevkit\images\train           （images\train 所在的绝对路径）

tarDir    = ——\VOCdevkit\images\val               (images\val 所在的绝对路径)

file_label_train   = ——\VOCdevkit\labels\train      （labels\train 所在的绝对路径）

file_label_val      =——\VOCdevkit\labels\val         （labels\val 所在的绝对路径）

运行代码即可得到训练集和验证集

训练参数设置：

1.在data文件夹下创建xxx.yaml ：可以复制随便一个yaml文件然后改个名字（我的为GKS.yaml）和内容

                                                         

内容如下：（具体看注释）

# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
# Train command: python train.py --data voc.yaml
# Default dataset location is next to /yolov5:
#   /parent_folder
#     /VOC
#     /yolov5


# download command/URL (optional)
#download: bash data/scripts/get_voc.sh

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: D:\yolov5\yolov5-5.0\yolov5-5.0\VOCdevkit\images\train  # 存放自己的训练集
val: D:\yolov5\yolov5-5.0\yolov5-5.0\VOCdevkit\images\val      # 存放自己的验证集

# number of classes
nc: 4        #自己的标签分类个数 在这里我的是四个  ↓

# class names
names: [ 'red','green','blue','bottle' ]  #标签分类名字 最好按照固定的顺序

然后进入如下目录：

                                             

复制yolov5s.yaml(根据自己所需模型大小选择)   改名为xxx.yaml（这里我的为yolov5_GKS.yaml）

其他都不动，修改nc：xxx（修改为你的标签分类数量，我的是4）

至此基本内容已经准备好了！

然后进入 train.py

 改一下这几个路径，对应上面自己定义的yaml文件夹即可

parser.add_argument("--weights", type=str, default="D:\yolov5\yolov5-master\weights\yolov5s.pt", help="initial weights path")  #权重文件需要自己下载
parser.add_argument("--cfg", type=str, default="D:\yolov5\yolov5-master\models\yolov5_GKS.yaml", help="model.yaml path")   #在models下的yaml文件
parser.add_argument("--data", type=str, default="D:\yolov5\yolov5-master\data\GKS.yaml", help="dataset.yaml path")          #在data下的yaml文件

其他参数可以看网上的自定义

至此可以开始训练了！

YOLOv8：

相比于YOLOv5 YOLOv8配置更简单一些：

XML格式换为txt，以及划分数据集验证集：

1.首先在yolov8目录下新建文件data                                                         D:\yolov8\data 
2.在data下创建images文件夹用于存放                                                     D:\yolov8\data\images (把训练图片放在这里面)
3.在data下创建Annotations文件夹用于存放打标签xml文件                         D:\yolov8\data\Annotations (把训练xml文件放在这里面)
4.在data下创建labels文件夹用于存放YOLO格式txt文件                     D:\yolov8\data\labels (暂时为空)

整体结构如下：

2.在data下创建：1.XML转TXT格式.py 文件 

复制黏贴以下代码进行XML格式转化为TXT格式：      （具体细节看注释）

import xml.etree.ElementTree as ET
import os, cv2
import numpy as np
from os import listdir
from os.path import join

classes = ['red','green','blue','bottle']  #修改自己的训练分类


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def convert_annotation(xmlpath, xmlname):
    with open(xmlpath, "r", encoding='utf-8') as in_file:
        txtname = xmlname[:-4] + '.txt'
        txtfile = os.path.join(txtpath, txtname)
        tree = ET.parse(in_file)
        root = tree.getroot()
        filename = root.find('filename')
        img = cv2.imdecode(np.fromfile('{}/{}.{}'.format(imgpath, xmlname[:-4], postfix), np.uint8), cv2.IMREAD_COLOR)
        h, w = img.shape[:2]
        res = []
        for obj in root.iter('object'):
            cls = obj.find('name').text
            if cls not in classes:
                classes.append(cls)
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            bb = convert((w, h), b)
            res.append(str(cls_id) + " " + " ".join([str(a) for a in bb]))
        if len(res) != 0:
            with open(txtfile, 'w+') as f:
                f.write('\n'.join(res))


if __name__ == "__main__":
    postfix = 'jpg'
    imgpath = 'D:\yolov8\data\images'          #修改为自己images的绝对路径
    xmlpath = 'D:\yolov8\data\Annotations'     #修改为自己Annotations的绝对路径
    txtpath = 'D:\yolov8\data\labels'          #修改为自己要储存的txt文件的绝对路径

    if not os.path.exists(txtpath):
        os.makedirs(txtpath, exist_ok=True)

    list = os.listdir(xmlpath)
    error_file_list = []
    for i in range(0, len(list)):
        try:
            path = os.path.join(xmlpath, list[i])
            if ('.xml' in path) or ('.XML' in path):
                convert_annotation(path, list[i])
                print(f'file {list[i]} convert success.')
            else:
                print(f'file {list[i]} is not xml format.')
        except Exception as e:
            print(f'file {list[i]} convert error.')
            print(f'error message:\n{e}')
            error_file_list.append(list[i])
    print(f'this file convert failure\n{error_file_list}')
    print(f'Dataset Classes:{classes}')

运行完后会发现labels文件夹下会多出txt文件，这就代表成功了！

2.在data下创建：2.划分测试训练验证集.py 文件

复制黏贴以下代码进行训练集测试集验证集的划分：

# coding=gb2312
import os
import random
import shutil

# 原数据集目录
root_dir = 'D:\yolov8\data'      #改为data文件夹的绝对路径
# 划分比例
train_ratio = 0.7
valid_ratio = 0.2
test_ratio = 0.1

# 设置随机种子
random.seed(42)

# 拆分后数据集目录
split_dir = 'D:\yolov8\data'     #改为data文件夹的绝对路径 （跟上面是同一个路径）
os.makedirs(os.path.join(split_dir, 'train', 'images'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'train', 'labels'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'valid', 'images'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'valid', 'labels'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'test', 'images'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'test', 'labels'), exist_ok=True)

# 获取图片文件列表
image_files = os.listdir(os.path.join(root_dir, 'images'))
label_files = os.listdir(os.path.join(root_dir, 'labels'))

# 随机打乱文件列表
combined_files = list(zip(image_files, label_files))
random.shuffle(combined_files)
image_files_shuffled, label_files_shuffled = zip(*combined_files)

# 根据比例计算划分的边界索引
train_bound = int(train_ratio * len(image_files_shuffled))
valid_bound = int((train_ratio + valid_ratio) * len(image_files_shuffled))

# 定义移动函数
def move_files(src_image_dir, src_label_dir, dest_image_dir, dest_label_dir, files):
    for image_file, label_file in files:
        src_image_path = os.path.join(src_image_dir, image_file)
        src_label_path = os.path.join(src_label_dir, label_file)
        dest_image_path = os.path.join(dest_image_dir, image_file)
        dest_label_path = os.path.join(dest_label_dir, label_file)
        shutil.copyfile(src_image_path, dest_image_path)
        shutil.copyfile(src_label_path, dest_label_path)

# 移动文件到划分后的目录
move_files(os.path.join(root_dir, 'images'), os.path.join(root_dir, 'labels'),
           os.path.join(split_dir, 'train', 'images'), os.path.join(split_dir, 'train', 'labels'),
           combined_files[:train_bound])

move_files(os.path.join(root_dir, 'images'), os.path.join(root_dir, 'labels'),
           os.path.join(split_dir, 'valid', 'images'), os.path.join(split_dir, 'valid', 'labels'),
           combined_files[train_bound:valid_bound])

move_files(os.path.join(root_dir, 'images'), os.path.join(root_dir, 'labels'),
           os.path.join(split_dir, 'test', 'images'), os.path.join(split_dir, 'test', 'labels'),
           combined_files[valid_bound:])

执行完后即可得到此次训练所需要的训练集、测试集、验证集了！

目录如下：

训练参数设置：

跟YOLOv5一样，本次我们需要准备两个yaml文件如图所示：

一个是xxx.yaml （我的叫data.yaml）训练集验证集测试集路径文件：

可以首先复制黏贴任意一个xxx.yaml文件到data目录下，然后删除所有内容再直接复制以下代码： 

# 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, ..]
train: D:\yolov8\data\train # 改为自己的训练集的绝对路径
val: D:\yolov8\data\valid # 改为自己的验证集的绝对路径
test: D:\yolov8\data\test # 改为自己的测试集的绝对路径

nc: 4  #改为训练类别个数

# Classes
names: ['red','green','blue','bottle']  #改为自己的训练类别名字

                                                                注意看注释

一个是xxx.yaml (我的叫yolov8n.yaml)训练集验证集测试集路径文件如下图所示：

可以复制这个路径下的文件（可选择自己的模型，我选的是yolov8.yaml）到data目录下，改为自己想要的名字，然后只需要修改nc：xxx为自己训练类别个数即可

然后我们还需要从官网上下载一个yolov8n.pt 文件到data目录下

官方网址如下：

https://github.com/ultralytics/ultralytics?tab=readme-ov-filehttps://github.com/ultralytics/ultralytics?tab=readme-ov-file点YOLOv8n 下载权重文件放到data目录下

至此我们的文件目录应该长这样：

然后如果我们要训练，直接在终端执行如下代码：

yolo task=detect mode=train model=D:\yolov8\data\yolov8n.yaml  data=D:\yolov8\data\data.yaml epochs=300 workers=5 device=0

其中model = yolov8n.yaml的绝对路径 data = data.yaml的绝对路径 epochs为训练轮数 device = 0代表使用GPU训练 如果没有可以去掉device = 0

至此就已经可以开始训练了！

附上摄像头检测代码：

import cv2
from ultralytics import YOLO

# 模型加载权重
model = YOLO(r'xxx\weights\best.pt')  #训练出最好的权重绝对路径
# 视频路径
cap = cv2.VideoCapture(0)
# 对视频中检测到目标画框标出来
while cap.isOpened():
    # Read a frame from the video
    success, frame = cap.read()

    if success:
        # Run YOLOv8 inference on the frame
        results = model(frame,conf=0.7)


        # Visualize the results on the frame
        annotated_frame = results[0].plot()

        # Display the annotated frame
        cv2.imshow("YOLOv8 Inference", annotated_frame)

        # Break the loop if 'q' is pressed
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        # Break the loop if the end of the video is reached
        break

# Release the video capture object and close the display window
cap.release()
cv2.destroyAllWindows()

如有不懂希望可以在评论区多多交流！谢谢！