用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