用yolov5训练自己的数据集

使用Yolov5v4.0训练皮带机煤炭识别数据集的详细教程

最新推荐文章于 2024-04-27 00:01:55 发布

原创最新推荐文章于 2024-04-27 00:01:55 发布 · 585 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #python

本文详细介绍了如何使用Yolov5v4.0版本对皮带机上的煤炭进行视频识别，包括环境配置、数据集构建（图片和标签制作、划分）、配置文件修改、训练过程中的报错解决以及测试和目标检测应用。

用yolov5训练自己的数据集

版本：yolov5 v4.0

一、环境配置（略）

二、数据集构建

1、本方案主要运用yolov5做皮带机上煤炭的视频识别检测，收集158张现场图片作为数据集，一共2个分类。（图略）

2、从GitHub中下载yolov5 v4.0版本文件并将其命名为belt_machine_detect，在belt_machine_detect中新建目录data，并新建images和labels文件夹，其中：

~images中存放所有的数据集158张图片（.jpg格式）

~labels中用于存放与数据集对应的标签文件（.txt格式）

3、标记数据集

先安装labelimg工具，在anaconda中进入yolov5虚拟环境中，输入pip install labelimg进行安装，安装完成后输入labelimg即可打开软件，在“Open Dir”中选择要标记的文件夹（data/images）打开，在“Change Save Dir”中选择标记生成的.txt文件保存路径，存放在data\labels文件夹中，注：labelimg输出格式选择YOLO。

4、划分数据集

默认train，val，test按照8：1：1的比例进行随机分类，本方案只划分train和val，在belt_machine_detect新建databases文件夹，在其中创建images和labels文件夹，这两个文件夹中都有train和val文件夹，将data中的数据集按照9:1的比例划分，并分别存放在train和val中。

5、修改配置文件

~5.1 在data文件夹中复制一份voc.yaml文件并重命名为voc_belt.yaml，打开并修改其中的内容如下：

train: databases/images/train # 训练集路径
val : databases/images/val  # 验证集路径
nc:2  # 类别数量
names : [ 'block', 'iron' ]  # 类别名称

~5.2 在models文件夹中复制一份yolov5s.yaml文件并重命名为yolov5s_belt.yaml，打开并修改其中的内容如下：

nc: 2 # 类别数量

三、训练数量集

在终端执行命令：

python train.py  --data data/voc_belt.yaml --weights weights/yolov5s.pt --epoch 100 --batch-size 4 --cfg models/yolov5s_belt.yaml

报错：

1、将database.py和general.py中的np.int改为np.int_
2、将loss.py中的：
# 原代码
anchors = self.anchors[i]
# 修改为
anchors,shape = self.anchors[i], p[i].shape
# 原代码
indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))  # image, anchor, grid indices
# 修改为
indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1)))

四、测试训练结果

python test.py  --data data/voc_belt.yaml --weights runs/train/exp8/weights/best.pt  --batch-size 4

五、正式目标检测（视频文件）

python detect.py  --source data/video/document2.mp4 --weights runs/train/exp8/weights/best.pt