在YOLOv5上训练自建数据集

目录

        本文记录的是在ubuntu18.06上使用YOLOv5（第六版）训练自建数据集的过程。

1.环境配置：

2.随机划分数据集：

3.训练yolov5s

---

1.环境配置：

        略（后续补充）。

2.随机划分数据集：

        首先将数据集整理成如下格式，其中images文件夹中保存的是划分成训练集和验证集的图片，labels文件夹中保存的是和images文件夹中对应图片的标签文件。

        其次在/data文件目录下创建和数据集对应的litchi.yaml文件，用作网络训练时加载数据的“指引文件”。nc表示类别，names表示各类别的名称。

path: ../yolov5-master/data/litchi_data  # dataset root dir yolov5-master/data/litchi_data
train: images/train  # train images (relative to 'path') 128 images
val: images/val  # val images (relative to 'path') 128 images
#test:  # test images (optional)

# Classes
nc: 1  # number of classes
names: ['y']  # class names

        下面是随机划分数据集的python代码，将数据集按照7:3的比例划分训练集和验证集，处理好数据之后就可以训练了。

import os
import random
from shutil import copy
#实现功能：将图片、标签文件随机划分并且创建符合训练规范的文件格式

#随机划分比例 训练集：验证集=7:3
train_rate=0.7
#原始图像及标签文件夹
image_Dir="/home/ice/Date_Stage_Y/Stage_Y_Img"
text_Dir="/home/ice/Date_Stage_Y/Stage_Y_Text"
#创建符合训练标准的文件夹
image_train="/home/ice/yolov5-master_v6/yolov5-master/data/litchi_data/images/train"
image_val="/home/ice/yolov5-master_v6/yolov5-master/data/litchi_data/images/val"

text_train="/home/ice/yolov5-master_v6/yolov5-master/data/litchi_data/labels/train"
text_val="/home/ice/yolov5-master_v6/yolov5-master/data/litchi_data/labels/val"

dir_list=[image_train,image_val,text_train,text_val]

for i in dir_list:
    if not os.path.exists(i):
        os.makedirs(i)

#获取标签目录下所有文件，并返回为列表格式
text_list=os.listdir(text_Dir)
sum_num=len(text_list)
print(text_list)

train_nums=int(train_rate*sum_num)
#返回[0,1,2,3,4,5...,sum_nums]
list=range(sum_num)
#随机一段长为train_nums的字符串,返回一个列表
train_index=random.sample(list,train_nums)
print(train_index)

for i in range(len(text_list)):
    # 如果为训练集，则将图像和标签复制到对应的文件夹
    if i in train_index:
        #图像
        copy(os.path.join(image_Dir,text_list[i].split(".")[0]+".jpg"),image_train)
        #标签
        copy(os.path.join(text_Dir,text_list[i]),text_train)
    #验证集
    else:
        #图像
        copy(os.path.join(image_Dir,text_list[i].split(".")[0]+".jpg"),image_val)
        #标签
        copy(os.path.join(text_Dir,text_list[i]),text_val)

3.训练yolov5s

        下载预训练权重有助于加速网络收敛，在YOLOv5的github项目上作者提供了v5各个版本预训练权重的下载（文件是以.pt结尾的）。在本次训练中，如下图所示--weights表示为yolov5s.pt的路径，--cfg表示为yolov5s网络结构的路径（注意要将该.yaml文件中的nc值设置成与litchi.yaml中的值一样），--data表示为数据集存储路径，--hyp表示超参数设置。