使用yolov4训练自己的数据集

环境：ubuntu16.04 cuda8.0 cudnn6.0.1 GT1070
1，GitHub：https://github.com/AlexeyAB/darknet下载

2，编译；
①修改makefile文件

GPU=1
CUDNN=1
CUDNN_HALF=0#这里如果显卡计算能力小于7.0，不需要改为1
OPENCV=1
AVX=0
OPENMP=1
LIBSO=1
’‘’
DEBUG=1
‘’‘
NVCC=/usr/local/cuda-8.0/bin/nvcc#改为自己的路径

②make
备注：常见错误均可网上搜到，无非是版本问题

3，下载预训练模型：GitHub上均有

4，简单测试：

./darknet detect cfg/yolov4.cfg yolov4.weights data/dog.jpg

测试效果：

Done! Loaded 162 layers from weights-file 
data/dog.jpg: Predicted in 67.411000 milli-seconds.
bicycle: 92%
dog: 98%
truck: 92%
pottedplant: 33%

5，训练自己数据集：
①收集数据集，这个根据自身项目要求进行收集，网上有绝大多数需要的数据集，下载归类。

②对于数据集的预处理：可以将众多图片名称改为VOC数据集的格式，即000**.jpg。修改代码如下：

import os
path=input('请输入文件路径(结尾加上/)：')       
#获取该目录下所有文件，存入列表中
f=os.listdir(path)
n=0
for i in f:
    #设置旧文件名（就是路径+文件名）
    oldname=path+f[n]
    #设置新文件名
    newname=path+'000'+str(n+1)+'.jpg'
    #用os模块中的rename方法对文件改名
    os.rename(oldname,newname)
    print(oldname,'======>',newname)
    n+=1

③对数据集进行标记：本人使用labelImg工具进行标记，附上GitHub链接：https://github.com/tzutalin/labelImg 非常简单的安装和操作。

④对数据集进行增广：一般常用翻转、剪切、旋转，根据自身项目需要增广至所需数据集数量。（代码较多，不便粘贴，需要的小伙伴儿可以私聊）

⑤将数据集改为VOC数据集格式：
一、数据集文件夹格式
VOC数据集文件夹格式一般为：

      ---VOC***（可以为当天日期本人VOC531）
          ----Annotations（里面为全部的标记完的.xml文件）
          ----ImageSets
              ----Main
          ----JPEGImages（里面为全部的图片）

二、将数据集进行分割
一般分为训练集、验证集、测试集，分割比例可以自己调整（本人使用9-1）代码如下

import os  
import random 
xmlfilepath=r'/home/nph/darknet-master/VOC531/Annotations'   #change xml path                        
saveBasePath=r"/home/nph/darknet-master/"                        #change base path
trainval_percent=0.9                                             #adjust trainval percentage    
train_percent=0.9                                              #adjust train percentage        
total_xml = os.listdir(xmlfilepath)
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)  
ftrainval = open(os.path.join(saveBasePath,'VOC531/ImageSets/Main/trainval.txt'), 'w')    #change your path
ftest = open(os.path.join(saveBasePath,'VOC531/ImageSets/Main/test.txt'), 'w')    #
ftrain = open(os.path.join(saveBasePath,'VOC531/ImageSets/Main/train.txt'), 'w')    #
fval = open(os.path.join(saveBasePath,'VOC531/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()  

需要改动的路径已经标记出来，运行之后在Main文件下会产生四个txt文件：分别为trainval.txt、train.txt、test.txt和val.txt。里面是对应所属类别图片的名称，也就是前面命名号的000***（没有后缀）如图：

这就是图片的标签。

三、提取标签和坐标信息。
这一步是将标签的信息和xml文件中标记的坐标信息提取出来，代码如下：

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

sets=[('531', 'train'), ('531', 'val'), ('531', 'test')] #改为自己保存的VOC***后的数字信息

classes = ["supervisor_white", "technicist_blue", "worker_yellow", "manager_red" ] #引号里为自己的数据集分类


def convert(size, box):
    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(year, image_id):
    in_file = open('VOC%s/Annotations/%s.xml'%(year, image_id))
    out_file = open('VOC%s/labels/%s.txt'%(year, 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 year, image_set in sets:
    if not os.path.exists('VOC%s/labels/'%(year)):
        os.makedirs('VOC%s/labels/'%(year))
    image_ids = open('VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('%s/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
        convert_annotation(year, image_id)
    list_file.close()

需要修改的地方都已标记，运行之后会在VOC***文件下产生一个labels文件，里面是txt文件，存放的是每个图片标记框的信息，如图：

在VOC***文件夹目录下会产生三个txt文件，一般命名为：***（VOC后面的数字）_train(test/val).txt，里面存放的是照片的路径，如图：

这三个txt文件可以放入VOC文件内，方便后面训练时填写地址。

至此，VOC数据集格式已经全部改好，接下来就是改动YOLOv4网络文件进行训练。

⑥修改网络文件：
主要修改cfg names data文件
将darknet-master/cfg/yolov4-custom.cfg、darknet-master/cfg/coco.data和darknet-master/data/coco.names文件复制到自己的VOC***文件下
一、cfg文件
修改cfg文件（本人改为531.cfg）

[net]
# Testing  #训练过程将test三行注释
#batch=1
#subdivisions=1
# Training
batch=32 #根据自己显卡内存量修改，一般为16,32,64
subdivisions=8 #根据上行改写，一般为4,8,16
width=416 #大小可以自己修改，一般是16的倍数
height=416
'''
learning_rate=0.001
burn_in=1000
max_batches = 12000 #一般是classess*2000，即训练类别*2000
policy=steps
steps=9600,10800 #max_batches的0.8和0.9倍
scales=.1,.1







[convolutional]
size=1
stride=1
pad=1
filters=24 #（classes+5）)3
activation=linear


[yolo]
mask = 6,7,8
anchors = 12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401
classes=3 #训练类别数目
num=9
jitter=.3
ignore_thresh = .7
truth_thresh = 1
random=1
scale_x_y = 1.05
iou_thresh=0.213
cls_normalizer=1.0
iou_normalizer=0.07
iou_loss=ciou
nms_kind=greedynms
beta_nms=0.6
max_delta=5
#3个yolo层都需要改，都在cfg文件最下层

二、修改data文件
打开531.data文件，修改如下：

classes= 3
train  = /home/u/qingxin/potgraduate/Fire/darknet-master/Fire/Multi/VOC531/531_train.txt #此处就是上边生成的训练的图片路径文件，改为自己的路径即可
valid  = /home/u/qingxin/potgraduate/Fire/darknet-master/Fire/Multi/VOC531/531_val.txt #验证图片路径文件，改为自己路径即可
names = /home/u/qingxin/potgraduate/Fire/darknet-master/Fire/Multi/VOC531/3fire.names #names文件路径
backup = /home/u/qingxin/potgraduate/Fire/darknet-master/Fire/Multi/VOC531/backup #训练过程中训练模型保存地址，可以自己修改

三、修改names文件
打开531.names文件，修改如下：

supervisor_white
technicist_blue
worker_yellow
manager_red

也就是自己数据集分类的名称

至此，网络文件修改完毕，可以进行训练。

⑦进行训练
在darknet-master路径下输入指令

./darknet detector train VOC531/531.data VOC531/531.cfg yolov4.conv.137 -gpus 1,2 -dont_show

分别是训练指令-data文件-cfg文件-预训练模型-指定GPU
可以添加日志

./darknet detector train VOC531/531.data VOC531/531.cfg yolov4.conv.137 -gpus 1,2 -dont_show 2>&1 | tee VOC531/531.log

开始训练～

测试待更新～