yolov5：训练自己的数据模型

Scs_space

已于 2024-08-01 10:29:38 修改

阅读量1.1k

点赞数 9

文章标签： YOLO

于 2024-07-29 16:10:44 首次发布

本文链接：https://blog.csdn.net/qq_41478219/article/details/140768803

版权

上一篇 yolov5的完整部署讲解了yolov5模型的部署和运行

实际生活中，我们需要的检索目标，很多并不在官方给出的模型之中。那我们就需要训练自己的数据模型。

这个训练过程并不复杂，主要是如下几个步骤：

1.理解机器训练的步骤是什么？

准备目标训练照片
把每个照片中的目标位置准确标记出来。
把标记目标的分类名称和位置信息转换成数字数据。
用图片库里的照片对比标记过的数据进行训练。

2.开始训练

首先，我们在工程的根文件夹下，新建一个文件夹用于储存训练相关文件（VOCData）。这部分借鉴于CSDN的一篇文章。Yolov5训练自己的数据集_CSDN博客。我以这边文章为基础，详解一下每一步。

2.1 收集训练照片，并完成分类标记。

首先，我们在VOCData文件夹下创建images和Annotations文件夹，这个后面我们会用到。

这里我们要用到一个工具 labelimg，通过下面命令安装这个工具：

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

安装好以后，我们在VOCData文件夹下创建一个predefined_classed.txt文件，在这个文件中，我们可以写入你自己定义的类别名称。每一个类别换一行。这里我指定以了一个分类 eagle.

在命令管理器里输入下面命令，进入labelimg界面。

labelimg predefined_classes.txt

把收集的训练图片全部输入在VOCData目录下创建的images文件夹。

待标注图片数据的路径文件夹，选择images文件夹

保存类别标签的路径文件夹，选择Annotations 文件夹

这个按键可以说明我们标注的标签为voc格式，点击可以换成yolo或者createML格式。（这里选择voc格式）

点击View，会出现如图红色框框中的选项。最好和我一样把勾勾勾上。

点击 Create RectBox，框选图片中的目标。保存名称为分类名。这里是 eagle，然后点击Next Image切换写一个图片进行标记。

关闭labelimg,我们可以在Annotations 文件夹里，看到新生成的.xml文件。

2.2 定义训练集

训练集就是把模型标签分配出哪些用于训练，哪些用于验证，哪些用于测试，把这些标签的信息分别储存在不同的.txt文件里。方便训练时候读取这些信息。train.txt（保存训练集）,val.txt（保存验证集）,test.txt（保存测试集）,trainval.txt(保存训练验证集）

这些动作手动操作速度太慢，可以写一个py程序直接生成对应文件。split_train_val.py（借鉴于

CSDN文章）,这个文件放在VOCData文件夹下，直接通过 python split_train_val.py 运行。

# coding:utf-8

import os
import random
import argparse

parser = argparse.ArgumentParser()
# xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下，注意以下为相对路径
parser.add_argument('--xml_path', default='Annotations', type=str, help='input xml label path')
# 数据集的划分，地址选择自己数据下的ImageSets/Main，注意以下为相对路径
parser.add_argument('--txt_path', default='ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0  # 训练集和验证集所占比例。 这里没有划分测试集
train_percent = 0.9  # 训练集所占比例，可自己进行调整
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

for i in list_index:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

程序运行完成后会生成 ImagesSets\Main 文件夹，在这个文件夹中会生成测试集、训练集、训练验证集、验证集，若要分配测试集，需要更改第 14、15 行的代码，即更改对应的比例）

提取 Yolov5训练所需的图片地址集和图片对应的目标信息标签集合。

在VOCData目录下创建程序 text_to_yolo.py 并运行，将XML格式转yolo_txt格式。

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd
 
sets = ['train', 'val', 'test']
classes = ["bottle"]  # 改为自己的类别
abs_path = os.getcwd()
print(abs_path)
 
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(image_id):
    in_file = open('D:/Yolov5/yolov5/VOCData/Annotations/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('D:/Yolov5/yolov5/VOCData/labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_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
        # 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))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
 
wd = getcwd()
for image_set in sets:
    # 这里是绝对路径，需要根据自己的情况修改
    if not os.path.exists('D:/Yolov5/yolov5/VOCData/labels/'):
        os.makedirs('D:/Yolov5/yolov5/VOCData/labels/')
    image_ids = open('D:/Yolov5/yolov5/VOCData/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
 
    if not os.path.exists('D:/Yolov5/yolov5/VOCData/dataSet_path/'):
        os.makedirs('D:/Yolov5/yolov5/VOCData/dataSet_path/')
 
    list_file = open('D:/Yolov5/yolov5/VOCData/dataSet_path/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write('D:/Yolov5/yolov5/VOCData/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

程序运行完成后，会生成 labels 文件夹和 dataSet_path 文件夹。

其中 labels 中为不同图像的标注文件。每个图像对应一个txt文件，文件每一行为一个目标的信息，分别为class, x_center, y_center, width, height，这种为 yolo_txt格式。

dataSet_path文件夹包含三个数据集的txt文件，train.txt等txt文件为划分后图像所在位置的路径，如train.txt就含有所有训练集图像的路径。

2.3 定义一个自己的训练参数文件

在 yolov5 目录下的 data 文件夹下 新建一个 myvoc.yaml文件

train：训练图片路径
val：验证图片路径
nc: 同时训练的类数量
names: 训练的类名称集合

train: D:/yolo/yolov5-7.0/VOCData/dataSet_path/train.txt
val: D:/yolo/yolov5-7.0/VOCData/dataSet_path/val.txt

# number of classes
nc: 1

# class names
names: ["eagle"]

2.4 设置训练文件参数

在yolov5目录下的model文件夹下是模型的配置文件，有n、s、m、l、x版本，逐渐增大（随着架构的增大，训练时间也是逐渐增大）。

打开 yolov5s.yaml，最好将yolov5s.yaml文件复制一份，然后将其重命名。将nc后面的参数改为训练的类数量，所有冒号后面需要加一个空格。

找到train.py这个py文件。找到主函数，这里面有模型的主要参数，训练自己的模型需要修改如下几个参数就可以训练了。

--weights：这个是网络模型选用yolov5s.pt。这个模型相对简单，训练速度快。准确度略低。

--cfg：默认yolov5s.yaml文件存放在models文件夹下，如果目录不同。请修改成自己的目录。

--data：这里需要设置为你自己定义的训练集参数文件。 myvoc.yaml

--hyp: 对应目录下如果有hyp.scratch-low.yaml，就设置这个。没有就设置hyp.scratch.yaml

--epochs：这个是训练的次数默认是300。单模型建议训练在300以上，少了训练出的模型识别不出目标。

--batch-size：这个是模型梯度。默认是16，数值越大，需要的内存和算力越大。配置不高，减小这个数值。

其他参数暂时不动。

2.5 启动训练

如果，你按照我的建议修改了参数。可以直接开始训练。

python train.py

也可以在执行train.py 后添加新的训练参数。

python train.py --weights weights/yolov5s.pt  --cfg models/yolov5s.yaml  --data data/myvoc.yaml --epoch 300 --batch-size 8 --img 640   --device cpu

训练好的模型会被保存在 yolov5 目录下的 runs/train/weights/ 下，生成 best.pt,last.pt文件。

3. 测试训练结果

这时我们打开detect.py这个文件，更改以下参数：

--weights：更改成新生成的best.pt

--source:更为测试图片存放地址。

执行 python detect.py，在runs文件夹下回生成，测试结果图片。

打开文件，你就可以看到你的训练成果了。

Scs_space

关注

9
点赞
踩
19

收藏

觉得还不错? 一键收藏
0
评论
复制链接

分享到 QQ

分享到新浪微博

扫一扫