YOLOv5训练自己的目标检测模型

YOLOv5 训练自己的目标检测模型

一、项目环境

操作系统 Windows11
Pytorch 1.12.1(基于CPU)
编程语言 Python3.10
anaconda版本 23.3.1

二、数据集制作

YOLOv5项目下载github地址：https://github.com/ultralytics/yolov5；下载之后文件名称为：yolov5-master。
在该文件夹中新建如下结构文件夹：

dataset（文件名可自定义）
–images
----train（训练集）
----val（验证集）
–labels
----train（训练集标签文件）
----val（存验证集标签文件）

标签文件和训练用的图片名称一一对应。

1、收集原始数据集图片

1)图片类型数据：可以直接使用
2)视频类型数据：利用opencv进行视频抽帧

import cv2

video=cv2.VideoCapture('yolo_test_blue.mp4')#打开对应视频
num=0#用于帧计数
gap=4#帧间隔，每隔4帧提取一张图片

while True:
    ret,frame=video.read()
    if not ret:
        print('video ends...')
        break
    num+=1
    if num%gap==0:
        #将图片存于"./yolov5-master/images/train"
        cv2.imwrite('./images/frame_{}.jpg'.format(int(num//gap)),frame)

2、用labelImg进行数据标注

1）labelimg下载
在命令窗使用pip install labelImg下载
2）labelimg使用
激活环境后直接输入labelimg即可打开

启动后界面如下：

点击①后会跳出两次路径选择，第一次是数据集的路径，第二次是标注集的保存路径；如果只跳出一次，可以点击②修改

我们将标注集保存在“./labels/train”中。

标注物体：

选中需要标注的物体，输入名称；

labelimg快捷键
W：创建选择方框
A：上一张图片
D：下一张图片

如果出现labelimg闪退的问题，可能是版本不兼容，可以使用不同的python版本安装labelimg，相当于创建不同版本python的环境，在该环境下重新安装labelimg。

>conda create -n envname python=3.9  #envname 环境名称
>conda activate envname  #激活环境
>pip install labelImg    #下载labelimg
>labelimg                #打开

标注完会生成classes.txt文本，里面保存的是我们定义的物体的名称。

将”./images/train”中的图片集选择一部分移动到“./images/val“中作为检验数据，将标注集中对应的部分标注集同样移动到”./labels/val“中；为了方便文件管理，将classes.txt移动到”./dataset“中。

三、模型训练

1）打开”.data/coco128.yaml“文件，修改对应参数：

path：文件夹的路径
train：训练集的路径
val：检验集路径
names：物体名称

修改为如下内容（我同时修改了文件名，只需与下一步保持一致即可）：

2）修改train.py参数

将’data’变量修改为刚才修改的.yaml文件的路径。

然后运行train.py开始训练模型；

运行结束后提示：

更新的文件夹如下：

四、目标检测

需要将 detect.py 中的 --weights参数修改为上一步生成weights文件夹中的 best.pt 或者 last.pt。

在命令窗中更改到yolov5-master路径下输入以下命令运行程序：

python detect.py --weights runs/exp6/weights/best.pt --source 0 --view-img
#打开摄像头的数据并显示结果