计算机视觉|YoloV5建立自己的数据集并进行训练 python|深度学习

本文记录本人基于yolov5训练夜间车辆识别模型的过程，数据采集自行车记录仪的视频片段，训练平台为Win11，python3.8，RTX3070 laptop
首先将视频片段按照10秒间隔切分，代码如下：

import cv2 import os import shutilpath = 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 = 300for 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 shutilTrainVal_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_numTxtList = 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盘，可以将它调整到空间较大的其他磁盘进行训练，步骤如下：
打开查看高级系统设置，点击高级-设置，打开性能选项，