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,界面如下

YoloV5建立自己的数据集并进行训练_第1张图片

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

YoloV5建立自己的数据集并进行训练_第2张图片

勾选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,修改数据集路径以及分类数目

YoloV5建立自己的数据集并进行训练_第3张图片

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

YoloV5建立自己的数据集并进行训练_第4张图片

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

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

YoloV5建立自己的数据集并进行训练_第5张图片

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

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

打开查看高级系统设置,点击高级-设置,打开性能选项,
YoloV5建立自己的数据集并进行训练_第6张图片
点击高级-更改,打开虚拟内存
YoloV5建立自己的数据集并进行训练_第7张图片
取消最上方自动管理勾选,选择自定义每个驱动器的分页文件大小,这里直接在磁盘E分配1T虚拟内存,完成自定义后点击设置,确定,关闭选项卡。而后重启电脑,生效设置。
YoloV5建立自己的数据集并进行训练_第8张图片

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

你可能感兴趣的:(计算机视觉,深度学习,python)