详细图文步骤，使用darknet框架和yolov4-tiny训练自定义数据集

前言

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

  • 测试结果
    

准备自定义数据集

• 收集进击的巨人相关的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”]

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
      

• 开始训练

  • 在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

成功！