前言#

之前讲了如何安装配置 YOLOv5 ，并使用默认的模型进行检测¹，这里开始尝试标注数据集，并使用数据集在 YOLOv5s 模型上进行训练和检测。

我的配置: Ubuntu 20.04 + Python 3.8

2022.1.8 找到了官方的教程: Train Custom Data

labelImg#

标注软件推荐使用 labelImg 。

Github: https://github.com/tzutalin/labelImg

作者没有提供 ubuntu 下的可执行文件，所以需要直接下源码运行。

执行如下指令，修改默认分类:

1
vim data/predefined_classes.txt

如果没有特殊需要，建议清空，否则标注时会默认增加这些分类，如果不注意会导致训练错误。

执行如下指令，配置运行环境:

1
sudo apt-get install pyqt5-dev-tools
2
sudo pip3 install -r requirements/requirements-linux-python3.txt
3
make qt5py3

执行如下指令，运行软件:

1
python3 labelImg.py
2
python3 labelImg.py [IMAGE_PATH] [PRE-DEFINED CLASS FILE]

下面是 labelImg 的界面，把红框里的数据格式改成 yolo，然后就能标注数据集了。

标注#

这里感谢 Vanessa Ni 的教程²，并提供了我的第一个数据集来源。

选择了一个鱼类相关的数据集: Labeled Fishes in the Wild

选择60张图片，将图片从0开始编号:

新建训练数据文件夹 images/ 和 labels/，将图片放到 images 里，然后使用 labelImg 标注，把标注结果保存到 labels 里:

完成数据集的标注后，创建数据集的配置文件。

将 images 和 labels 都放在 YOLO/traindata/ 文件夹下，在该路径下创建 traindata.yaml 文件，用于配置数据集信息:

1
# train and val data
2
train: /home/nono/YOLO/traindata/
3
val: /home/nono/YOLO/traindata/
4

5
# number of class
6
nv: 1
7

8
# class names
9
names: ['fish']

• train : 训练数据集的路径
• val : 测试数据集的路径
• nv : 分类数量
• name : 分类名词

训练#

进入 YOLO 项目的根目录。

输入如下指令，使用 YOLOv5s 对数据集进行训练:

1
python train.py --weights yolov5s.pt --data traindata/traindata.yaml --img 320 --nosave --device cpu

• weights : initial weights path，初始权重文件
• data : hyperparameters path，训练数据文件
• img : train, val image size (pixels)，训练输入图像的大小
• nosave : only save final checkpoint，只保存最后和最好的权重文件
• device : cuda device, i.e. 0 or 0,1,2,3 or cpu，选择计算单元

详细的可通过 python train.py -h 查看帮助

最后的训练结果如下图:

用了59张图片，共计130个标签，最后 mAP 达到了 0.992。

最终的训练结果能在 runs/train/ 找到。

下图是 YOLO 给的关于模型性能的分析图表:

下图是模型在训练集上的检测结果:

检测#

进入 YOLO 项目的根目录。

输入如下指令，用新训练的模型进行检测:

1
python detect.py --weights runs/train/exp10/weights/best.pt --source traindata/test.mp4 --device cpu

• weights : model path(s)
• source : file/dir/URL/glob, 0 for webcam
• device : cuda device, i.e. 0 or 0,1,2,3 or cpu

详细的可以通过 python detect.py -h 查看帮助

检测结果如下图:

最终的检测结果和视频在 runs/detect/ 中。

报错#