YoloV4训练自己的数据集

---

1-建立工作文件夹

新建一个项目文件夹，用于存放接下来训练所需要的文件。

mkdir train    # 或者其他英文名
1

切换至工作文件夹

cd train
1

按照下图所示建立相关文件夹和文件。

.
├── JPEGImages
├── Annotations
├── labels
├── backup
├── data
│   ├── train.data
│   ├── train.names
│   ├── yolov4.cfg
│   ├── yolov4-tiny.cfg
├── darknet
├── gen_files.py
├── yolov4.conv.137
├── yolov4-tiny.conv.29
1
2
3
4
5
6
7
8
9
10
11
12
13
14

相关文件夹的作用 相关文件的作用

gen_files.py内容如下：

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random

# 类别列表
classes=["ball"]


# 递归清除Linux隐藏文件
def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)


# PASCAL VOC格式的xml标注文件 转换为 YOLO格式的txt文本标注文件
def convert(size, box):
    '''
    size = (w, h)
    box = (xmin, xmax, ymin, ymax)
    '''
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)


# 执行单个标注文件的转换
def convert_annotation(image_id):
    in_file = open('Annotations/%s.xml' %image_id)
    out_file = open('labels/%s.txt' %image_id, 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    in_file.close()
    out_file.close()


# 当前目录
wd = os.getcwd()

# 检查是否存在Annotations文件夹
annotation_dir = os.path.join(wd, "Annotations/")
if not os.path.isdir(annotation_dir):
        os.mkdir(annotation_dir)
# 清除隐藏文件
clear_hidden_files(annotation_dir)

# 检查是否存在JPEGImages文件夹
image_dir = os.path.join(wd, "JPEGImages/")
if not os.path.isdir(image_dir):
        os.mkdir(image_dir)
# 清除隐藏文件
clear_hidden_files(image_dir)

# 检查是否存在backup文件夹
backup_dir = os.path.join(wd, "backup/")
if not os.path.isdir(backup_dir):
        os.mkdir(backup_dir)
# 清除隐藏文件
clear_hidden_files(backup_dir)

# 检查是否存在labels文件夹
labels_dir = os.path.join(wd, "labels/")
if not os.path.isdir(labels_dir):
        os.mkdir(labels_dir)
# 清除隐藏文件
clear_hidden_files(labels_dir)

# 新建文件train.txt、test.txt
# 存放需要训练和测试的完整文件路径
train_file = open(os.path.join(wd, "train.txt"), 'w')
test_file = open(os.path.join(wd, "test.txt"), 'w')
train_file.close()
test_file.close()


# 训练数据集
train_file = open(os.path.join(wd, "train.txt"), 'a')
# 测试数据集
test_file = open(os.path.join(wd, "test.txt"), 'a')

# 列出所有图片文件
list = os.listdir(image_dir)
# 设置训练集/测试集划分比例的随机数
probo = random.randint(1, 100)
print("Probobility: %d" % probo)
for i in range(0,len(list)):
    path = os.path.join(image_dir,list[i])
    if os.path.isfile(path):
        image_path = image_dir + list[i]
        # 根据文件名，得到没有后缀的文件名和后缀名
        (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
        # 标注文件名
        annotation_name = nameWithoutExtention + '.xml'
        # 标注文件地址
        annotation_path = os.path.join(annotation_dir, annotation_name)
    # 设置训练集/测试集划分比例的随机数
    probo = random.randint(1, 100)
    print("Probobility: %d" % probo)
    # 训练集和测试集的划分，这里的75代表训练集/测试集的划分比例为75:25
    # 训练集
    if(probo < 75):
        if os.path.exists(annotation_path):
            # 在当前目录下的train.txt文本文件中，写入训练图片的完整地址
            train_file.write(image_path + '\n')
            # 执行标注格式转换
            convert_annotation(nameWithoutExtention)
    # 测试集
    else:
        if os.path.exists(annotation_path):
            # 在当前目录下的test.txt文本文件中，写入训练图片的完整地址
            test_file.write(image_path + '\n')
            # 执行标注格式转换
            convert_annotation(nameWithoutExtention)
# 文件操作结束后，关闭文件流
train_file.close()
test_file.close()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

---

2-准备训练数据集

2-1 将训练所需的图片拷贝至JPEGImages文件夹下。

2-2 将训练图片对应的XML标注文件拷贝至Annotations下。

确保labels文件夹下没有文件及隐藏文件
确保训练图片和标注一一对应

2-3 修改gen_files.py中的classes为自己的标签。

例如：

classes = ["person", "phone", "chair"]
1

有几个标签修改为几个。

2-4 在终端运行gen_files.py

python3 gen_files.py
1

此时会在当前训练文件夹生成文本文件train.txt和test.txt。
labels文件夹下会生成YOLO格式的txt标注文件。

train.txt文件内容为训练集图片的绝对地址的集合，一行一条。
test.txt文件内容为测试集图片的绝对地址的集合，一行一条。

可以看到train.txt和test.txt中的条目数的比例大致为之前设置的75:25。
当然也可以设置为其他的训练集/测试集划分比例。

labels下的文件是JPEGImages文件夹下每一个图像的YOLO格式的标注文件,这是由Annotations文件夹的xml标注文件转换来的。

最终训练只需要:

• train.txt
• test.txt
• labels文件夹下的txt文本标注文件
• JPEGImages文件夹下的图像文件

此时训练文件夹的目录树如下所示：

.
├── JPEGImages
├── Annotations
├── labels
├── backup
├── data
│   ├── train.data
│   ├── train.names
│   ├── yolov4.cfg
│   ├── yolov4-tiny.cfg
├── darknet
├── gen_files.py
├── yolov4.conv.137
├── yolov4-tiny.conv.29
├── train.txt
├── test.txt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

新增train.txt、test.txt。
labels文件夹下新增n个YOLO格式的txt标注文件。

---

3 修改配置文件

3-1 新建data/train.names文件

可以复制darknet目录下的data/voc.names到训练目录下的data/train.names。
再根据自己情况的修改，可以重新命名如：data/train.names。

names文件存放训练的标签，一行一个标签，不要有空行。

例如3个标签的情况下，可以修改为：

person
phone
chair
1
2
3

替换为自己训练的标签即可。

3-2 新建data/train.data文件
可以复制darknet目录下的cfg/coco.data到训练目录下的data/train.data。

再根据自己情况的修改，可以重新命名如：data/train.data。

data文件存放类别数、训练集位置、测试集位置、names文件位置、训练出来的模型的存放地址等信息。

例如：

classes = 80
train  = train.txt
valid  = test.txt
#valid = data/coco_val_5k.list
names = data/train.names
backup = backup
eval = coco
1
2
3
4
5
6
7

这里的文件地址都是相对于训练工作文件夹的。

train.data相关参数及其作用

---

3-3和3-4二选一，看需要训练的模型是YoloV4还是YoloV4-Tiny。

---

3-3 新建data/yolov4-tiny.cfg
可以复制darknet目录下的cfg/yolov4-tiny.cfg到训练目录下的yolov4-tiny.cfg。
再根据自己情况的修改，可以重新命名data/yolov4-tiny-xxx.cfg。

batch = 64
subdivisions = 32
1
2

在data/yolov4-tiny.cfg文件中，两个yolo层和各自前面的convolutional层的参数需要修改:

• 两个yolo层都要改：yolo层中的classes为类别数
• 每一个yolo层前的convolutional层中的filters =(类别+5)* 3

例如:
yolo层：classes=1；convolutional层：filters=18
yolo层：classes=2； convolutional层：filters=21
yolo层：classes=4；convolutional层：filters=27

3-4 新建data/yolov4.cfg

可以复制darknet目录下的cfg/yolov4.cfg到训练目录下的yolov4.cfg。
再根据自己情况的修改，可以重新命名data/yolov4-xxx.cfg。

batch = 64
subdivisions = 32
1
2

在data/yolov4.cfg文件中，两个yolo层和各自前面的convolutional层的参数需要修改:

• 两个yolo层都要改：yolo层中的classes为类别数
• 每一个yolo层前的convolutional层中的filters =(类别+5)* 3

例如:
yolo层：classes=1；convolutional层：filters=18
yolo层：classes=2； convolutional层：filters=21
yolo层：classes=4；convolutional层：filters=27

---

4-Anchor Box先验框聚类分析与修改

4-1 使用k-means聚类获得自己数据集的先验框大小

对于YoloV4-Tiny:

./darknet detector calc_anchors data/train.data -num_of_clusters 6 -width 416 -height 416
1

对于YoloV4:

./darknet detector calc_anchors data/train.data -num_of_clusters 9 -width 416 -height 416
1

4-2 修改cfg文件中的先验框大小

cfg文件中的anchors位置的几个数字

4-3 重新训练和测试

5-训练自己的数据集

5-1 确保yolov4-tiny.conv.29或者yolov4.conv.137在训练文件夹下

5-2 开始训练

训练命令：

./darknet detector train data文件地址 cfg文件地址 预训练权重文件(不含yolo head层) -map
1

如需无须显示训练过程的map变化,在命令末尾加-map，即：

./darknet detector train data/train.data data/yolov4-tiny.cfg yolov4-tiny.conv.29
1

如需要显示训练过程的map变化,在命令末尾加-map，即：

./darknet detector train data/train.data data/yolov4-tiny.cfg yolov4-tiny.conv.29 -map
1

对于YoloV4模型的训练，只需要把命令里的cfg配置文件和预训练权重文件替换为YoloV4版本的就可以。

5-3 训练建议

• batch=64
• subdivisions=4(或2,1)
• YOLOv4: 把max_batches设置为 (classes*2000);但最小为4000。 YOLOv4-tiny可减少
• YOLOv4: 把steps改为max_batches的80% and 90%;例如steps=3200, 3600。YOLOv4-tiny可相应减少
• 为增加网络分辨率可增大height和width的值,但必须是32的倍数 (height=608, width=608或者是32的整数倍) 。这有助于提高检测精度

---

6-测试训练出来的网络模型

训练好后可以在backup目录下看到权重文件。

测试图片

./darknet detector test data/train.data data/yolov4-tiny.cfg backup/yolov4-tiny_best.weights xxx.jpg
1

测试视频

./darknet detector demo data/train.data data/yolov4-tiny.cfg backup/yolov4-tiny_best.weights xxx.mp4
1

---

7-性能统计

模型的性能主要体现在：mAP

6-1 统计 mAP@IoU=0.50:

./darknet detector map data/train.data data/yolov4-tiny.cfg backup/yolov4-tiny_best.weights
1

6-2 统计 mAP@IoU=0.75:

./darknet detector map data/train.data data/yolov4-tiny.cfg backup/yolov4-tiny_best.weights -iou_thresh 0.75
1