YoloV5建立自己的数据集并进行训练

本文记录本人基于yolov5训练夜间车辆识别模型的过程，数据采集自行车记录仪的视频片段，训练平台为Win11，python3.8，RTX3070 laptop

首先将视频片段按照10秒间隔切分，代码如下：

import cv2
import os
import shutil

path = r'D:\yolo\imageprocess\movie'
SavePath = r'D:\yolo\imageprocess\image'
files = os.listdir(path)

isExists = os.path.exists(SavePath)
if not isExists:
    os.makedirs(SavePath)
    print('path of %s is build' % SavePath)
else:
    shutil.rmtree(SavePath)
    os.makedirs(SavePath)
    print('path of %s already exist and rebuild' % SavePath)

count = 300

for movie in files:
    if movie.endswith('.MP4'):
        MoviePath = os.path.join(path, movie)

        VidCap = cv2.VideoCapture(MoviePath)
        fps = int(VidCap.get(cv2.CAP_PROP_FPS))  # 获取视频每秒的帧数

        i = 0
        j = 0
        while True:
            success, image = VidCap.read()
            i += 1
            if i % count == 0:
                j += 1
                SaveName = movie.split('.')[0] + '_' + str(j) + '_' + str(i) + '.jpg'
                cv2.imwrite(SavePath + '/' + SaveName, image)
            if not success:
                print('video %s is all read' % (movie.split('.')[0]))
                break

而后运行labelImg，对图像进行标注：

首先进入labelImg-master\data\predefined_classes.txt文件，对labelImg预定义的目标分类进行修改，方便后续快速标注

激活yolov5环境，进入labelImg文件夹

conda activate yolov5py38
cd /d D:\yolo\labelImg-master\requirements         //labelImg-master依赖包路径
pip install -r requirements-linux-python3.txt         //安装依赖包
cd /d D:\yolo\labelImg-master
python labelImg.py

运行labelImg，界面如下

选择Open Dir，依次选择图片所在路径与Label保存路径

勾选Use default label，方便快捷选择类型
主界面左侧可以选择切换标注类型，VOC或Yolo，这里改选为Yolo
在view菜单中勾选Auto Save Mode激活自动保存

标注过程中常用的快捷键如下

w                         开始创建矩形框；
d                         切换到下一张图；
a                         切换到上一张图；
del                      删除选中的标注矩形框；
Ctrl++                 放大图片；
Ctrl--                   缩小图片；
↑→↓←                 移动选中的矩形框的位置
Ctrl + Shift + d         删除当前图片；

其他常见用快捷键如下

Ctrl + u                 选择要标注的文件目录；
Ctrl + r                  选择标注好的标签存放的目录；
Ctrl + s                 保存标注好的标签（自动保存模式下会自动保存）；
Ctrl + d                 复制当前标签和矩形框；
Ctrl + Shift + d      删除当前图片；
Space                   将当前图像标记为已验证；

完成所有标注后，将数据集随机分割为训练集与验证集

​
import os
import random
import shutil

TrainVal_percent = 0.1
Train_percent = 0.9
ImageFilePath = r'D:\yolo\imageprocess\all_images'
TxTFilePath = r'D:\yolo\imageprocess\all_labels'
total_txt = os.listdir(TxTFilePath)
num = len(total_txt)
Train_num = int(num * Train_percent)
Val_num = num - Train_num

TxtList = range(num)
ValList = random.sample(TxtList, Val_num)

Train_image_path = r'D:\yolo\imageprocess\images\train'
Val_image_path = r'D:\yolo\imageprocess\images\val'

Train_txt_path = r'D:\yolo\imageprocess\labels\train'
Val_txt_path = r'D:\yolo\imageprocess\labels\val'

if not os.path.exists(Train_image_path):
    os.makedirs(Train_image_path)
if not os.path.exists(Val_image_path):
    os.makedirs(Val_image_path)
if not os.path.exists(Train_txt_path):
    os.makedirs(Train_txt_path)
if not os.path.exists(Val_txt_path):
    os.makedirs(Val_txt_path)

for i in TxtList:
    name = os.path.splitext(total_txt[i])[0]
    name_txt = name + '.txt'
    name_jpg = name + '.jpg'
    if i in ValList:
        shutil.copy(os.path.join(TxTFilePath, name_txt), os.path.join(Val_txt_path, name_txt))
        shutil.copy(os.path.join(ImageFilePath, name_jpg), os.path.join(Val_image_path, name_jpg))
    else:
        shutil.copy(os.path.join(TxTFilePath, name_txt), os.path.join(Train_txt_path, name_txt))
        shutil.copy(os.path.join(ImageFilePath, name_jpg), os.path.join(Train_image_path, name_jpg))

分割后的文件结构如下：

├─mydataset
│  ├─images
│  │  ├─train
│  │  └─val
│  └─labels
│      ├─train
│      └─val
└─yolov5-master

在yolov5-master/data文件夹下新建一个mydataset.yaml文件，参考coco128.yaml，修改数据集路径以及分类数目

修改模型的配置文件，选择yolov5s作为预训练模型，进入...\yolov5-master\models，修改yolov5s.yaml，只需要修改检测目标类型个数（yolov5会自动生成anchors）

而后下载预训练模型权重文件yolov5s.pt，正常情况下运行train.py时会自动下载权重文件，如果速度太慢可以去github下载，链接如下yolov5s.pt

下载完成后将yolov5s.pt放在yolov5-master根目录下，而后打开train.py，对以下参数进行修改。

修改完成后，运行train.py即可开始训练。

训练开始后遇到了页面分页太小的问题，提示OSError: [WinError 1455] 页面文件太小,无法完成操作。这是由于系统分配的虚拟内存大小不够，数据集过大会溢出，需要进入系统设置，重新分配虚拟内存。一般情况下虚拟内存被分配在系统盘C盘，可以将它调整到空间较大的其他磁盘进行训练，步骤如下：

打开查看高级系统设置，点击高级-设置，打开性能选项，

点击高级-更改，打开虚拟内存

取消最上方自动管理勾选，选择自定义每个驱动器的分页文件大小，这里直接在磁盘E分配1T虚拟内存，完成自定义后点击设置，确定，关闭选项卡。而后重启电脑，生效设置。

完成虚拟内存的重新分配后，即可运行train.py进行训练。