前言
darknet框架优点是易于安装、没有依赖项。可以很方便的训练yolov3、yolov4、yolov3-tiny、yolov4-tiny等多种yolo系列网络。
本文使用yolov4-tiny网络,收集了一个小小的自定义数据集,识别动漫人脸,主要在unbuntu18.04系统下完成。
主要工作
可执行的darknet文件 通过配置、编译Makefile生成
自定义数据集 配置.data文件
配置网络参数 修改yolov4-tiny.cfg文件中的filters和classes
正文
配置darknet
- 下载darknet框架,地址:https://github.com/pjreddie/darknet ,解压后进入darknet-master文件夹
- 配置Makefile文件
-
使用
gedit Makefile
,打开Makefile文件
参数说明
GPU=1 #使用GPU设置为1,不使用GPU设置为0
CUDNN=1 # 不使用设置为0
CUDNN_HALF=0
OPENCV=0 #如果要配置OpenCV则设置为1
-
使用
make -j8
编译Makefile文件 ,编译成功的话会在当前目录下生成可执行的darknet文件
注意事项
第一次编译时直接调用make -j8
,如果要重新编译Makefile,要先用make clean
清除旧的darknet,再重新编译可能出现的问题
因为cudnn版本的问题,编译时可能会出现错误: fatal error: cudnn.h: No such file or directory
解决方案:https://github.com/pjreddie/darknet/issues/2356 -
测试darknet
- 在YOLOv4 model zoo中下载yolov4.weights,放到当前目录下
- 输入下述命令进行测试
./darknet detector test ./cfg/coco.data ./cfg/yolov4.cfg ./yolov4.weights data/dog.jpg -i 0 -thresh 0.25
- 或者直接运行
bash image_yolov4.sh
- 在YOLOv4 model zoo中下载yolov4.weights,放到当前目录下
-
测试结果
-
准备自定义数据集
- 收集进击的巨人相关的jpg图像20张,任务是检测其中的人脸
- 数据存放目录如下。主目录为attackgiant,其下JPEGImages用来存放图片,Annotations用来存放即将标注的标签
- 使用LabelImg进行数据格式标注
- Open Dir选择JPEGImages
- Change Save Dir选择Annotations。
- 数据格式建议选择voc格式,如果后续还想用其它网络训练,也比较方便
- 因为数据集只有一个face标签,所以勾选 Use default label
- 标注快捷键
- 画框,w
- 保存, ctrl+s
- 下一张,d
- 上一张,a
- 数据配置
- 进入build/darknet/x64/文件夹,使用该文件夹存放自定义的数据并进行后续操作,与源码区分离开
- 在build/darknet/x64/下新建datasets文件夹,将attackgiant文件夹拷贝到datasets下
- 在attackgiant下新建prepare_dataset.py,代码如下,用于生成yolo格式的标注文件
- 修改dataset_path为了数据集的绝对路径,如
dataset_path='/databank/home/user/darknet-master/build/darknet/x64/datasets/attackgiant/'
- train_percent为训练集的百分比,train_percent=0.8表示80%用于训练,本文就是20*0.8=16幅图像用于训练
- classes为所有的标签名,如classes = [“face”]
- 修改dataset_path为了数据集的绝对路径,如
import os
import random
import xml.etree.ElementTree as ET
import cv2
#数据集的绝对路径
dataset_path='/databank/home/user/darknet-master/build/darknet/x64/datasets/attackgiant/'
xmlfilepath='./Annotations'
saveBasePath="./ImageSets"
trainval_percent=1
#80%用于训练集
train_percent=0.8
temp_xml = os.listdir(xmlfilepath)
total_xml = []
#所有标签名
classes = ["face"]
for xml in temp_xml:
if xml.endswith(".xml"):
total_xml.append(xml)
num=len(total_xml)
list=range(num)
tv=int(num*trainval_percent)
tr=int(tv*train_percent)
trainval=random.sample(list,tv)
train=random.sample(trainval,tr)
print("train and val size",tv)
print("traub suze",tr)
if not os.path.exists(saveBasePath):
os.makedirs(saveBasePath)
ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')
ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')
ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')
fval = open(os.path.join(saveBasePath,'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()
def convert_annotation(image_id,part_name):
in_file = open(dataset_path+'Annotations/%s.xml' % image_id,'r', encoding='UTF-8')
out_file_img = open(dataset_path+part_name, 'a')
if not os.path.exists(dataset_path+'labels/'):
os.makedirs(dataset_path+'labels/')
out_file_label = open(dataset_path+'labels/%s.txt' % image_id,'a')
tree = ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
voc_img_dir=dataset_path+'JPEGImages/{}.jpg'.format(image_id)
out_file_img.write(voc_img_dir)
out_file_img.write("\n")
img=cv2.imread(voc_img_dir)
dh = 1. / img.shape[0]
dw = 1. / img.shape[1]
cnt=len(root.findall('object'))
if cnt==0:
print(image_id)
cc=0
for obj in root.iter('object'):
cc+=1
cls = obj.find('name').text
if cls not in classes:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
if dw*float(xmlbox.find('xmin').text)<0. or dw*float(xmlbox.find('xmax').text)<0. or dh*float(xmlbox.find('ymin').text)<0. or dh*float(xmlbox.find('ymax').text)<0.:
print(image_id)
b = (dw*float(xmlbox.find('xmin').text), dw*float(xmlbox.find('xmax').text), dh*float(xmlbox.find('ymin').text),
dh*float(xmlbox.find('ymax').text))
out_file_label.write(str(cls_id)+ " " + str((b[0]+b[1])/2) + " " + str((b[2]+b[3])/2) + " " + str(b[1]-b[0]) + " " + str(b[3]-b[2]))
if cc<cnt:
out_file_label.write("\n")
out_file_label.close()
def generate_image_path(part_name):
imgname_list = []
with open(os.path.join(dataset_path, 'ImageSets/' + part_name)) as f:
all_lines = f.readlines()
for a_line in all_lines:
imgname_list.append(a_line.split()[0].strip())
print(len(imgname_list))
for image_id in imgname_list:
print(image_id)
convert_annotation(image_id,part_name)
generate_image_path('train.txt')
generate_image_path('val.txt')
- 运行成功后的目录如下
-
ImageSets存放训练集、测试集的图像名
-
labels是yolo格式的标签
-
train.txt是训练集的path
-
- 新建classes.names,存放标签名
- 新建attackgiant.data,train和valid分别是刚才生成的训练集和测试集path,backup用来存放训练时的权重文件,此处的路径均是相对于build/darknet/x64/路径的
训练配置
-
将编译生成的darknet文件拷贝到 build/darknet/x64/下
-
修改cfg文件参数
- 在cfg目录下复制一份yolov4-tiny.cfg文件,重命名为yolov4-tiny-attackgiant.cfg
- 修改yolov4-tiny-attackgiant.cfg
- 修改文件中两处的filters和classes。filters的修改规则为 (classes + 5) * 3
classes是对每种类别进行预测的概率,+5是预测每个box的x、y、w、h和置信度5个值,乘3是因为每个网格会进行3种尺度的预测。本文的种类只有1种,所以filters = (1 + 5) * 3 = 18
- 因为数据集很小很小,设置max_batches = 3000
- 修改文件中两处的filters和classes。filters的修改规则为 (classes + 5) * 3
-
开始训练
-
在build/darknet/x64/下输入下列命令开始训练
./darknet detector train datasets/attackgiant/attackgiant.data cfg/yolov4-tiny-attackgiant.cfg -map
-
如果放到后台训练,选择输出日志文件,也可以
nohup ./darknet detector train datasets/attackgiant/attackgiant.data cfg/yolov4-tiny-attackgiant.cfg -map > train-attackgiant.log 2>&1 &
-
-
进行测试
./darknet detector test datasets/attackgiant/attackgiant.data cfg/yolov4-tiny-attackgiant.cfg backup/yolov4-tiny-attackgiant_last.weights datasets/attackgiant/JPEGImages/img01.jpg