Darknet Yolov4训练自己的数据集
背景
环境：Ubuntu18.04+python3.6+显卡1080Ti+CUDA10.0+cudnn7.5.1+OpenCV3.4.6，Yolov4模型（入门级）
注意：这里的python最好用3.6版本的，3.7版本的python环境执行python_images.py有点版本不兼容的小问题.(个人碰到的)
文章从上往下，步骤是依次运行的
Darknet工程
从github克隆下载源码，链接地址：https://github.com/AlexeyAB/darknet
cd 此工程的根目录
编译
1.修改makefile文件里面的一些参数值（根据自己实际情况来修改）
GPU=1
CUDNN=1
CUDNN_HALF=1
OPENCV=1
AVX=0
OPENMP=1
LIBSO=1
ZED_CAMERA=0
ZED_CAMERA_v2_8=0
2.用makefile编译
make -j8
此时你编译可能会报链接lcuda不成功的错误，解决方法如下：
I modified this :
ifeq ($(GPU), 1)
COMMON+= -DGPU -I/usr/local/cuda/include/
CFLAGS+= -DGPU
ifeq ($(OS),Darwin) #MAC
LDFLAGS+= -L/usr/local/cuda/lib -lcuda -lcudart -lcublas -lcurand
else
LDFLAGS+= -L/usr/local/cuda/lib64 -lcuda -lcudart -lcublas -lcurand
endif
endif
to :
ifeq ($(GPU), 1)
COMMON+= -DGPU -I/usr/local/cuda-8.0/include/
CFLAGS+= -DGPU
LDFLAGS+= -L/usr/local/cuda-8.0/lib64 -lcudart -lcublas -lcurand
LDFLAGS+= -L/usr/local/cuda-8.0/lib64/stubs -lcuda
endif
正常情况下你的darknet工程就编译成功了，此时你的根目录会多出一个darknet的可执行文件。
测试工程是否编译成功
下载预训练权重：（这里用到的是yolov4）
yolov4.cfg文件：https://raw.githubusercontent.com/AlexeyAB/darknet/master/cfg/yolov4.cfg
yolov4.weights文件：https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights
然后执行下面代码测试coco数据集：
# 测试图片，结果保存在darknet-master/predictions.jpg
./darknet detect cfg/yolov4.cfg yolov4.weights data/dog.jpg
训练自己的数据集
1.数据集创建
数据集目录结构是这个样子的：
darknet-master
---- data
---- obj.names # 物体类别名称(如果有两类物体，就写上两类物体的名称)
---- obj.data # 将数据集的信息保存在这个文件中，yolov4从这个文件中读取数据集信息
---- obj # 存放图片以及每个图片的标签信息
---- train.txt # 存放训练集地址 (相对地址,比如: data/obj/image1.jpg)
---- test.txt # 存在测试集地址 (相对地址,比如: data/obj/image3.jpg)
我分别介绍一个data文件夹下五个文件或文件夹的作用以及格式：
(1) obj.names 该文件中保存的是检测物体的名称。假设有一个安全帽检测的项目，检测有没有戴安全帽。没有戴安全帽的标签是people，戴安全帽的标签是hat。那么obj.names文件中就这样写：
person
hat
每一类名称独占一行，这样是为了方便读取文件中的内容，我们只需要通过换行符就可以轻松分割并读取obj.names中每个类别的名称。
(2) obj.data 该文件中保存着五类信息：类别数量，训练集，验证集，类别名称和保存权重的文件
classes= 2 # 2表示数据集中只有两类可检测的物体
train = data/train.txt # 表示保存训练数据集的地址
test= data/test.txt # 表示保存验证数据集的地址
names = data/obj.names # 表示可检测物体的名称
backup = backup/ # 表示保存训练权重的文件
obj.data文件其实就是一个汇总文件，yolov4需要的数据集地址，数据集标签以及信息的时候就是从这个文件中得到的。该文件涉及的内容在上下文都有涉及，这里就不再赘述。
(3) obj 该文件夹中存放着整个数据集（训练集和验证集）的图片（.jpg格式）以及它们的标签（.txt）文件。也就是说每一张图片都对应一个txt标签文件。比如image1.jpg图片对应的标签文件就是image1.txt文件。将数据集图片和标签放在一起有一个好处就是，我们不需要单独将图片和标签放在不同的文件夹下，只需要提供一个文件路径，得到所有的.jpg图片之后将其后缀改成.txt就可以得到相应的标签文件，这样做方便有简洁。
每个标签文件中包含如下五个信息：
<object-class> <x_center> <y_center> <width> <height>
object-class: 表示物体的数字标签。比如0, 1, 2等整数。
举例来说：
=
=
如果一幅图片中包含不止一个物体，那么每幅图片的标签信息应该如何填写？举例来说，对于image1.jpg图片有三个物体，image1.txt文件就是这样：
1 0.716797 0.395833 0.216406 0.147222
0 0.687109 0.379167 0.255469 0.158333
1 0.420312 0.395833 0.140625 0.166667
(4) train.txt 该文件中保存着所有训练集的相对地址。（可以是绝对路径）比如：
data/obj/img1.jpg
data/obj/img2.jpg
data/obj/img3.jpg
(5) test.txt 该文件中保存着所有验证集的相对地址。（可以是绝对路径）比如：
data/obj/img11.jpg
data/obj/img21.jpg
data/obj/img31.jpg
生成train.txt/test.txt的python脚本如下：
import os
import io
import math
import sys
import random
import argparse
from collections import namedtuple, OrderedDict
label_names = ['person','car','bus','truck']
def get_files(dir, suffix):
res = []
for root, directory, files in os.walk(dir):
for filename in files:
name, suf = os.path.splitext(filename)
if suf == suffix:
#res.append(filename)
res.append(os.path.join(root, filename))
return res
def gbbox_iou(box1, box2):
b1_x1, b1_y1, b1_x2, b1_y2 = box1
b2_x1, b2_y1, b2_x2, b2_y2 = box2
inter_rect_x1 = max(b1_x1, b2_x1)
inter_rect_y1 = max(b1_y1, b2_y1)
inter_rect_x2 = min(b1_x2, b2_x2)
inter_rect_y2 = min(b1_y2, b2_y2)
inter_width = inter_rect_x2 - inter_rect_x1 + 1
inter_height = inter_rect_y2 - inter_rect_y1 + 1
if inter_width > 0 and inter_height > 0:
inter_area = inter_width * inter_height
#iou
b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)
#iou = inter_area / (b1_area + b2_area - inter_area)
iou = inter_area / b1_area
else:
iou = 0
return iou
def convert_dataset(list_path, output_file):
# 读取目录里面所有的 txt标记文件 列表
label_list = get_files(list_path, '.txt')
total_label_len = len(label_list)
random.shuffle(label_list)
print('total_label_len', total_label_len)
error_count = 0
fp=open(output_file,'w')
for i in range(0, total_label_len):
sys.stdout.write('\r>> Calculating {}/{} error{}'.format(
i + 1, total_label_len, error_count))
sys.stdout.flush()
# 单个Label txt文件读取
label_file = label_list[i]
file_name, type_name = os.path.splitext(label_file)
image_path = file_name + '.jpg'
if type_name != '.txt' or not os.path.exists(image_path):
error_count += 1
print("error_file: ",label_file.encode('UTF-8', 'ignore').decode('UTF-8'))
continue
fd = open(label_file, 'r')
lines = [line.split() for line in fd]
fd.close()
error_id = 0
for line in lines:
class_index = int(line[0])
xmins = float(line[1]) - float(line[3]) / 2
ymins = float(line[2]) - float(line[4]) / 2
xmaxs = float(line[1]) + float(line[3]) / 2
ymaxs = float(line[2]) + float(line[4]) / 2
if float(line[3])<=0 or float(line[4]) <= 0 :
error_id = 1
print('\n error index: ', class_index, 'label_file', label_file)
continue
if class_index >= 3:
error_id = 1
print('\n error index: ', class_index, 'label_file', label_file)
continue
# if xmins < 0 or ymins < 0 :
# error_id = 1
# print('\n error index: ', class_index, 'label_file', label_file)
# if ymaxs > 1 or xmaxs > 1 :
# print('\n error index: ', class_index, 'label_file', label_file)
# error_id = 1
if error_id:
continue
# is_person_car = False
# bbox_num = len(lines)
# for i in range(0, bbox_num):
# if int(lines[i][0]) != 0:
# continue
# for j in range(0, bbox_num):
# if i==j or int(lines[j][0])==0:
# continue
# xmins = float(lines[i][1]) - float(lines[i][3]) / 2
# ymins = float(lines[i][2]) - float(lines[i][4]) / 2
# xmaxs = float(lines[i][1]) + float(lines[i][3]) / 2
# ymaxs = float(lines[i][2]) + float(lines[i][4]) / 2
# xmins1 = float(lines[j][1]) - float(lines[j][3]) / 2
# ymins1 = float(lines[j][2]) - float(lines[j][4]) / 2
# xmaxs1 = float(lines[j][1]) + float(lines[j][3]) / 2
# ymaxs1 = float(lines[j][2]) + float(lines[j][4]) / 2
# box1 = (xmins, ymins, xmaxs, ymaxs)
# box2 = (xmins1, ymins1, xmaxs1, ymaxs1)
# #过滤行人在车中
# iou = gbbox_iou(box1, box2)
# if iou > 0.99:
# is_person_car = True
# if is_person_car:
# continue
print("image_path: ", image_path)
fp.write(image_path)
fp.write('\n')
print('total_label_len', total_label_len)
fp.close()
def main():
parser = argparse.ArgumentParser(prog='gen_label_list.py')
parser.add_argument('--img-path', type=str, default='/root/zhangsong/fairworks/github/darknet-master/fireworks/data/smoke', help='test path')
parser.add_argument('--valid', type=str, default='fireworks/data/test_train.txt', help='*.txt path')
opt = parser.parse_args()
print(opt.img_path, opt.valid)
convert_dataset(opt.img_path, opt.valid)
if __name__ == '__main__':
main()
到这，数据集的搭建就已经完成了。下面我们需要对yolov4的配置文件进行修改。
3 修改配置文件
(1) 复制yolov4.cfg文件将其重命名为yolov4-obj.cfg。这样就不会对原始文件进行修改。
(2) 修改yolov4-obj.cfg如下内容：
# step1: 修改batch和subdivisions
L2: batch=64 # 原来就是64,根据gpu自己选择
L3: subdivisions=16 # 原来是8，根据自己的gpu选择
# step2: 修改图片的尺寸
L7: width=608 # 这边我就不进行修改了
L8: height=608 # 这边我也不修改
# step3: 修改classes(每个yolo层都需要修改一次，一共需要修改三次)
L968: classes=2 # 我只需要识别两类物体，因此需要修改成2
L1056: classes=2
L1144: classes=2
# step4: 需要修改每个yolo相邻的上一个convolution层的filter
L961: filters=21 # 因为我预测两类物体：21 = 3*(5+2)
L1049: filters=21
L1137: filters=21
2.进行训练
(1) 下载预训练权重: yolov4.conv.137
(2) 执行下面命令：
./darknet detector train data/obj.data cfg/yolov4-obj.cfg yolov4.conv.137 -map
(3) 在训练的过程中，darknet会自动保存权重文件
backup文件夹下的yolo-obj_last.weights文件会每隔100个iterations保存一次，新的会替代旧的
backup文件夹下的yolo-obj_xxxx.weights文件会每隔1000个iterations保存一次
darknet总的训练步长可以在yolov4-obj.cfg文件中修改max_batches的大小(默认max_batches = 500500)
补充：如果想要指定具体的gpu进行训练，可以使用-i来指定，比如我想使用索引为2的gpu进行训练，可以这样写：（*.weights文件也可以做预训练模型）
./darknet detector train data/obj.data cfg/yolov4-obj.cfg yolov4.conv.137 -i 2 -map
或者保存log日志
./darknet detector train fireworks/firesmoke.data fireworks/firesmoke.cfg fireworks/firesmoke.conv.23 -map 1> fireworks/train_smoke.log
测试自己训练的模型
训练好权重之后，在终端中输入下面命令就行：
测试图片：
./darknet detector test ./cfg/coco.data ./cfg/yolov4.cfg ./yolov4.weights data/dog.jpg -i 0 -thresh 0.25
测试视频：
./darknet detector demo ./cfg/coco.data ./cfg/yolov4.cfg ./yolov4.weights test50.mp4 -i 0 -thresh 0.25
测试网络摄像头：
./darknet detector demo ./cfg/coco.data ./cfg/yolov3.cfg ./yolov3.weights rtsp://admin:[email protected]:554 -i 0 -thresh 0.25
用自己训练的权重作为预训练
有时候训练到一半突然终止了，这时候从头开始训练又很费时间，此时我们可以将自己之前保存的权重作为预训练权重。但是直接使用yolo-obj_last.weights会报错。需要做出如下转变。
#首先用第一行代码将yolo-obj_last.weights转化为olo-obj_last.conv.23
./darknet partial cfg/yolo-obj.cfg backup/yolo-obj_last.weights backup/yolo-obj_last.conv.23 23
#第二行将我们刚转化好的yolo-obj_last.conv.23作为预训练权重训练
./darknet detector train cfg/obj.data cfg/yolo-obj.cfg backup/yolo-obj_last.conv.23