概述
yolo在目标检测领域是非常有代表性的模型,它速度快识别效果也很精准,是实时检测模型中应用最广泛的。yolo的原理和代码是很容易获得的,且有各式各样的教程,但是模型怎么使用的教程相对比较少。本文讲解如何使用yolov5模型训练自己的模型,同时也包含了如何标注自己的数据。
标注自己的数据
首先要收集自己的要训练的数据,可以通过网络收集或者自己拍摄的途径获得。收集好图片之后就要进行图片的标注,这里推荐使用labelimg标注工具标注,可以自行在网络搜索获得。安装完成labelimg标注软件之后,打卡标注软件,软件内部如下图所示。
在标注之前我们要将搜集的图片存放在一个文件夹下,然后新建一个文件夹用来存放标注文件。点击软件左上角的"Open Dir"图标打开自己存放图片的路径,接着点击"Change Save Dir"图标打开刚刚创建的保存标注文件的路径。然后点击下图箭头指向的地方,点击切换为"YOLO"图标(代表使用yolo格式存储,便于后续训练)。
完成上述操作之后就可以开始标注图片了。按w新建标注框,编辑标签名称,保存后就成功标注了第一张图片。
保存后查看标注文件保存的路径会看到新建了两个文件,classes.txt代表类别,另一个代表刚刚标注的文件对应的类别和坐标信息和原图片名称相对应。
在代码中配置自己的标注文件
标注就演示到这里了。这里只演示了单张的图片的标注,实际的项目中样本量至少也要几十张,多张标注方式和这个一样。
标注完成数据之后要划分训练数据集和验证数据集,可以按8:2的比例划分。划分时候的目录结构可以参考下图。
标注完数据之后就要进行训练了,yolov5使用的是ultralytics版本的代码,可以自行在github中下载。打卡工程文件之后找到下图所示的data路径,新建一个自己的yaml文件,
打开新建的yaml文件,将自己的图片数据和对应标注文件划分好训练和验证的路径填写在“train”和“val”后,代码如下:
train: D:\mydata\train# 训练数据的路径
val: D:\mydata\val# 验证数据的路径
# Classes
names:
0: fire #这里填写自己标注的类别名称。
训练自己的数据集
经过以上操作数据集的标注和配置文件的编写就完成了,接下来要进行到训练步骤了。当然在训练模型之前要配置好自己的代码环境,在项目根目录下可以找到requirements.txt文件,根据里面的对应的软件包名称完成环境配置。配置好环境之后,我们找到yolov5项目文件的根目录,在终端中打开。在终端中输入以下命令,–data后加自己的创建的yaml文件,–weight代表预训练权重 使用yolov5s.pt文件即可,–cfg代表模型配置文件使用yolov5s.yaml即可,–batch-size 代表同时训练的图片张数(如果显卡的显存较小,适当写小一点,在硬件支持的条件下越大越好),–epochs代表训练的轮次。
python train.py --data mydata.yaml --weight yolov5s.pt --cfg yolov5s.yaml --batch-size 10 --epochs 100
如没有问题就开始训练了,终端会输出训练的进度。
训练结束之后终端里会输出生成的pt文件存放路径,以及这次训练的精度的统计,还有存放这次训练的结果路径,可以通过这些数据回顾整个训练过程。保存的pt文件可以为自己模型推理和预测。
用训练的模型pt文件进行预测
训练完自己的数据集之后得到了pt文件,这个文件可以预测自己的目标检测对象,预测在yolov5中使用的detect.py文件,和上文一样在终端中打开项目文件的根目录,在终端中输入以下命令。–weight使用自己的pt文件,–source 代表用来预测的图片或视频的路径。
python detect.py --weight fire.pt --source D:\video\fire.mp4
执行上述代码之后会进行预测,在对应的结果目录找到预测的结果便能知晓模型效果。下图为火焰检测的效果。