【YOLO】YOLOv4训练自己的数据集
虽然在笔者写这篇博客的时候，YOLOv5已经问世，但是看到许多博客都说YOLOv4效果更佳，具体如何，笔者得亲测才能认同。
YOLOv4论文链接
YOLOv4代码链接
1 环境搭建
1. 创建虚拟环境并激活
conda create -n yolo4 pip python=3.6
conda activate yolo4
2. Requirements 安装必要的依赖项
根据github：
conda install cudatoolkit=10.0
conda install cudnn=7.6.5
conda install cmake # 这里笔者安装的cmake=3.18.2
conda install opencv # 这里笔者安装的opencv=3.4.2笔者选择的框架是tensorflow，如果电脑允许最好下载GPU版本的tensorflow
根据github上的指示，进入tensorflow框架下的yolov4源码github的链接：
根据指示继续安装requirements
pip install --ignore-installed --upgrade tensorflow-gpu==2.3.0 -i https://pypi.douban.com/simple
pip install opencv-python==4.1.1.26; lxml; tqdm...3. 下载源代码
git clone https://github.com/AlexeyAB/darknet.git编译
由于后续训练要使用GPU，因此需要修改makefile文件，具体修改部分如图所示：
根据上图的标注部分进行修改
darknet YOLO 编译使用GPU
- NVCC要对应自己的电脑路径进行修改
要根据电脑的计算能力修改，笔者的电脑无法满足，于是将其注释掉,选择可以满足的计算能力，可以百度搜索自己电脑的GPU 计算能力。
在终端运行，编译：
1 # cd到darknet文件夹下：
2 make # 或make -j82 下载开源权重并测试
yolov4.weights官网链接
笔者百度网盘链接：https://pan.baidu.com/s/1h4vuNFgWRyiMoppNnVR1wA
提取码：h0c5
./darknet/results detect cfg/yolov4.cfg yolov4.weights data/dog.jpg其中，./darknet/results结果保存路径；data/dog.jpg测试图片
3 训练自己的数据集
数据集结构如下图所示：
其中，Annotations: 标签文件（.xml）
ImageSets/Main: 在三个.txt中存放三类数据集的图片名称
将train_val_test.py代码复制到ImageSets相同根目录下，生成 需要的.txt文件，代码如下：
可以自行修改数据集的比例。
"""
2020.09.18:alian
divide train&val&test save as txt
"""
import os
import random
from pathlib import Path
trainval_percent = 0.8 # trainval数据集占所有数据的比例
train_percent = 0.75 # train数据集占trainval数据的比例
xmlfilepath = 'Annotations' # 标注数据xml储存的路径
txtsavepath = 'ImageSets/Main' # txtb保存的路径
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
print('total number is ', num)
list = range(num)
tv = int(num * trainval_percent)
print('trainVal number is ', tv)
tr = int(tv * train_percent)
print('train number is ', tr)
print('test number is ', num - tv)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
# windows
# ftrainval = open(txtsavepath+'\ '+'trainval.txt', 'w')
# ftest = open(txtsavepath+'\ '+'test.txt', 'w')
# ftrain = open(txtsavepath+'\ '+'train.txt', 'w')
# fval = open(txtsavepath+'\ '+'val.txt', 'w')
# ubuntu
ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/Main/val.txt', 'w')
for i in list:
name = total_xml[i][:-4] + '\n'
if i in trainval:
ftrainval.write(name)
if i in train:
ftrain.write(name)
else:
fval.write(name)
else:
ftest.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()JepgImages: 原始图片
1. 准备YOLOv4需要的label和txt
./darknet/scripts下的voc_label.py复制到项目根目录下./darknet，并对其内容进行修改:
"""
2020.11.09
author:alian
"""
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
# sets=[('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
sets=[('turnout', 'train'), ('turnout', 'val'), ('turnout', 'test')] # 以类别名称命名
# classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
classes = ["turnout"] # 类别名称
def convert(size, box):
dw = 1./(size[0])
dh = 1./(size[1])
x = (box[0] + box[1])/2.0 - 1
y = (box[2] + box[3])/2.0 - 1
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(classes, image_id):
# in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
# out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, image_id), 'w')
in_file = open('VOC2007_%s/Annotations/%s.xml'%(classes, image_id))
out_file = open('VOC2007_%s/labels/%s.txt'%(classes, 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')
wd = getcwd()
for classes, image_set in sets:
if not os.path.exists('VOC2007_%s/labels/'%(classes)):
os.makedirs('VOC2007_%s/labels/'%(classes))
image_ids = open('VOC2007_%s/ImageSets/Main/%s.txt'%(classes, image_set)).read().strip().split()
list_file = open('%s_%s.txt'%(classes, image_set), 'w')
for image_id in image_ids:
list_file.write('%s/VOC2007_%s/JPEGImages/%s.jpg\n'%(wd, classes, image_id))
convert_annotation(classes, image_id)
list_file.close()
# os.system("cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt > train.txt")
# os.system("cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt > train.all.txt")
os.system("cat turnout_train.txt turnout_val.txt > train.txt")
os.system("cat turnout_train.txt turnout_val.txt turnout_test.txt > train.all.txt")执行如下命名：
在根目录下将会生成训练需要的文件，即各个训练集中包含图像的路径。
cd ./darknet
python voc_label.py
2、修改配置文件
（1）. cfg目录下：（voc.data/coco.data二选一即可，笔者这里选择cov.data）
- 类别数量
- 训练过程中训练数量和验证数量的txt文件（即上述voc_label.py生成的文件）
- 类别标签名称
- 存放权重的路径
（2）yolov-obj.cfg
yolov4训练参数和相关网络结构的修改 ：
复制cfg文件夹下的yolov4-custom.cfg重命名为yolo-obj.cfg，修改如下参数：
输入图像大小和训练测试阶段中的batch的数量和划分次数；
图像大小必须是32的倍数
训练代数 ：
github中给出了max_batches的基本设置方法，2000 ×classes（但不少于训练图像的总数以及6000）。当然，设置的大一些也是可以的，只不过后期基本上在某一值附近震荡
值得注意的是，steps的设置是max_batches × 80% 和 max_batches × 90%。
网络结构 ：
根据待测目标类别的数量更改YOLO层(3处地方)和YOLO层前一层的卷积层（3处地方）
- 包含YOLO前一层卷积层的卷积核个数：（classes + 5）*3 （三处）
- YOLO层的类别数classes （三处）
- 锚框（可选，kmeans聚类） （暂时不改）
（3） data目录下：
开始训练
下载预训练权重：yolov4.conv.137
笔者百度云盘链接：https://pan.baidu.com/s/1p-7zbIIDNqQrOZGAWG3B6Q
提取码：zoiw
放到./darknet文件夹下面
- 训练指令：
/darknet detector train cfg/voc.data cfg/yolo-obj.cfg yolov4.conv.137 -gpus 6,7 # 多GPU训练训练好的权重文件将会保存在./darknet/backup文件夹下面
- 若在服务器上训练，无法可视化损失窗口，则使用如下命令：
/darknet detector train cfg/voc.data cfg/yolo-obj.cfg yolov4.conv.137 -dont_show- 在没有GUI的远程服务器上进行训练期间查看mAP和损失图，使用命令
./darknet detector train cfg/voc.data cfg/yolo-obj.cfg yolov4.conv.137 -dont_show -mjpeg_port 8090 -map在谷歌或者火狐浏览器上查看http://ip-address:8090
- 继续训练指令：
./darknet detector train cfg/voc.data cfg/yolo-obj.cfg backup\yolo-obj_2000.weights其中：backup\yolo-obj_2000.weights为最后的模型
预测
1 ./darknet detector test cfg/voc.data cfg/yolov4-custom.cfg yolov4-custom_xxxx.weights