YOLO v5 训练自定义数据集
Darknet有两个版本，分别对应着v3和v4；在v4后不久，就有了pytorch版本的v5，据称是性能和效果更好。这里主要是翻译如何用yolo v5来训练自定义的数据集，加了些自己尝试的过程，参考了三篇文章，链接见文章末尾。
下载yolov5的github代码：
git clone https://github.com/ultralytics/yolov5
安装所需的环境：
pip install -U -r yolov5/requirements.txt
如果之后遇到了matplotlib中的错'raise ImportError("Failed to import any qt binding")'，那么运行pip install PyQt5可以解决；
然后需要准备自己的图像数据集，并标记好类别和位置，可以用labelImg来帮助自己生成yolo数据(https://github.com/tzutalin/labelImg)，效果如下图：
点击左边的create rectbox来勾选范围，并输入类别，记得将类别由pacalvoc切换为yolo以免再次转换，然后点击保存，就会在指定目录生成一个class.txt文件包含类别信息，以及另一个txt文件包含与该图片有关的类别区域信息，如下所示：
# class x_center y_center width height
1 0.337500 0.408626 0.320522 0.131579
1 0.340112 0.541667 0.327239 0.108187
0 0.340112 0.694444 0.330970 0.125731
由于自己做数据会比较久，可以下载一份公共平台上的血细胞检测数据集来尝试，地址是：https://public.roboflow.com/object-detection/bccd，大概8M左右，这里需要下载版本为'YOLO v5 pytorch'。
下载后解压，放在yolov5的目录中
- yolov5
- data
- models
- trainDatasets # 解压到这里
- test
- images
- labels
- train
- images
- labels
- valid
- images
- labels
- data.yaml
- utils
- weights
将data.yaml中修改下路径：
train: ./trainDataset/train/images
val: ./trainDataset/valid/images
test: ./trainDataset/test/images
nc: 3
names: ['Platelets', 'RBC', 'WBC']
在models文件夹下创建一个custom_yolov5.yaml，实际是对应着yolov5s这个model，少许修改：
# parameters
nc: {num_classes} # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
# anchors
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Focus, [64, 3]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, BottleneckCSP, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 9, BottleneckCSP, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, BottleneckCSP, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 1, SPP, [1024, [5, 9, 13]]],
[-1, 3, BottleneckCSP, [1024, False]], # 9
]
# YOLOv5 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, BottleneckCSP, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
可以开始训练了，用cpu会慢一些，有gpu会快些
python train.py --img 416 --batch 16 --epochs 100 --data './trainDataset/data.yaml' --cfg ./models/custom_yolov5s.yaml --weights '' --name yolov5s_results --cache
- img: 代表输入的图像大小
- batch: 代表批量的大小
- epochs: 代表训练的周期，类似迭代，3000+都非常常见
- data: 指定数据集的配置信息
- cfg: 指定训练用的model配置信息
- weights: 指定路径来保存生成的weights文件
- name: 代表结果文件夹的名字
- nosave: 代表值保存最后的检查点
- cache: 代表缓存图像用以更快地训练
训练的过程的输出类似下图
因为图片并不多，于是训练时间不会多长，几个小时完全足够，训练完成后能看到生成了文件夹runs/yolov5s_results，里面有很多结果文件包含图表，在其中的weights文件夹中会有best.pt和last.pt文件，best.pt包含了最好的表现的weights，于是可以用它来做对应数据集的检测。
需要注意的是，pytorch的weights文件无法直接被opencv的cv2.dnn.readNetFromDarknet加载，所以这是不如v4的一个地方，需要转换为ONNX model然后用cv2.dnn.readNetFromONNX来加载。还好v5提供了转换的脚本，可参考https://github.com/ultralytics/yolov5/issues/239.
python detect.py --source trainDataset/test/images/BloodImage_00038_jpg.rf.63d04b5c9db95f32fa7669f72e4903ca.jpg --weights runs/train/yolov5s_results3/weights/best.pt --conf 0.25
运行这个命令便能通过训练后的weights文件来检测图片中的类别，其中conf代表着置信度阈值(thresholding objectness confidence)，会在runs下生成detect文件夹，里面便存放了检测结果，打开就能看到原图片基础上标记了识别出来的目标对象。
参考文章：
https://blog.roboflow.com/how-to-train-yolov5-on-a-custom-dataset/