本文章以瓶盖检测为例,教大家使用yolov5训练自己的数据集模型,并使用该模型进行目标检测.

在进行本文学习时应该准备好以下内容:

1. 完成anaconda的安装 链接: 超详细版Anaconda的安装及使用conda创建、运行虚拟环境以及使用镜像源
2. 完成数据集的标注和划分 链接: yolo模型训练所需的数据集的划分教程
3. 可以成功运行yolov5的目标检测(本文也会教) 链接: 使用YOLOv5模型进行目标检测
4. 安装CUDA与CUDNN(没有独显的同学可以跳过) 链接: CUDA,CUDNN的安装教程,让你可以使用Pytorch GPU显卡来跑程序

如果已经可以成功运行上面第三步目标检测,则可以跳过第一步,从2.1大步开始

1.虚拟环境的配置

1.1虚拟环境的创建

在win+R中打开cmd
输入以下代码创建一个名为yolov5s,python3.8版本的虚拟环境
名称和版本可以更改,但版本不要超过3.9,否则很多库都不兼容

conda create -n yolov5s python=3.8

等待之后出现以下输出则表示以及创建好,如果没创建好,可能是anaconda没用安装配置好,可以看完之前的文章 链接: 超详细版Anaconda的安装及使用conda创建、运行虚拟环境以及使用镜像源

输入以下代码进入到环境中

conda activate yolov5s

1.2所需库的安装

我们找到yolo的文件夹,复制文件夹的路径
在cmd中 cd 此路径

cd C:\Users\Apzzed\Desktop\AI\yolov5-master

然后输入以下代码,即可自动下载安装所有需要的库

pip install -r requirements.txt

如果下载速度过慢,可能会导致报错,可以复制下面代码(中科大源)到cmd中运行,然后在重新进行上面下载库的步骤

conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/cloud/msys2/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/cloud/bioconda/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/cloud/menpo/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/cloud/

耐心等待安装完成,如果出现黄色警告不要紧,出现红色才表示有错误

1.3 放入yolov5s.pt模型文件

把yolov5s.pt文件放入yolo文件夹的根目录下

1.4测试目标检测是否能正常运行

输入下面代码运行
source后的0代表用默认摄像头,也可以改为其他图片的路径
weights 后的文件表示用yolov5s.pt模型处理(要保证第一步中yolov5s.pt放在文件夹的根目录下)

python detect.py --source 0 --weights yolov5s.pt

也可以用pycharm打开detect.py,翻到最下面找到这两行,更改参数再运行,注意用pycharm运行时注意右下角环境选择好你刚才创建的环境

想要关闭时,如果再cmd 中运行,按ctrl+c 强制关闭;pycharm中按ctrl+F2或停止按钮关闭

至此环境配置已经完成

2.配置数据集路径文件

在data文件夹下.找到coco128.yaml文件,将他复制并粘贴到同目录下,并给他改一个名字,我把名字更改成cap.yaml并在pycharm中找到这个文件打开

如下图所示,下面五个部分都需要更改

1. path代表你的yolo的根目录路径
2. train代表你的训练集路径
3. val代表你的验证集路径,这俩个都是之前数据集划分后得来的文件夹
   *注意路径都要手动改成右斜杠才行,并注意在最后有没有再加一个右斜杠
4. nc表示你一共有多少个标注类型的数量
5. 下面name中表示你的各个标注类型的名称,这里我只有一个标注名为cap的

3.配置模型文件

在models文件夹中有一个yolov5s.yaml的文件,打开它

这里只需要更改标注类型的数量即可,数量与上一步保持一致

4.修改训练文件参数

打开yolo根目录下的train.py文件

往下滑到四百多行,找到这些代码需要更改
第一行时训练的权重文件,我们选中之前放进来的yolov5s.pt文件
第二行时是第三步中我们更改的文件的路径
第三行是第二布中我们更改的文件的路径
下面epochs指的是你要训练多少代,如果你是用cpu跑的,建议10代就好,速度快一些,当然越多代最终效果越好

再往下翻这个works该小一点,不然容易爆内存错误,0也可以

5.开始训练及成果

运行train.py文件

当最后出现了saved什么路径就是跑成功了,有其他报错不要紧,在这个路径下就可以找到你的成果

你的成果包括一些训练情况的分析图形,以及weights里有两个pt模型文件,是你最终训练出来的模型,分别是最好的一次和最后一次

6.使用自己训练的模型进行目标检测

在回到detect.py文件,将这两个更改
第一行改成刚才你训练出来的两个模型中的其中一个
第二行更改为随便一个测试用的图片
运行此文件

最后运行输出保存在这个路径下

我们找到这个路径,就可以看到检测的成果