YOLOv4训练自定义数据集

首先需要编译，cd到darknet-master目录下，

更改makefile文件

GPU=1

CUDNN=1

CUDNN_HALF=1

OPENCV=1 ……

make -j2

需要改动的文件为 darknet/src 下的image.c文件

找到draw_detections_v3 函数，在第369行处：

 qsort(selected_detections, selected_detections_num, sizeof(*selected_detections), compare_by_probs);
    for (i = 0; i < selected_detections_num; ++i) {

在for (i = 0; i < selected_detections_num; ++i) { 这个for循环中加入代码：

如果只需要一类识别结果，加入代码为：

if(strcmp(names[selected_detections[i].best_class], "你所需要检测的类别名字") != 0){
					continue;	
			}

如果需要两类识别结果，在if中再嵌套一个if条件：

if(strcmp(names[selected_detections[i].best_class], "你所需要检测的第一种类别名字") != 0){
			 	if(strcmp(names[selected_detections[i].best_class], "你所需要检测的第二种类别名字") != 0){
					continue;
				}	
			}

（不改这个部分好像也没什么问题，只是看到有一篇博客说要改这里T_T）

然后配置训练文件

先标注训练图片，用https://github.com/tzutalin/labelImg

建文件夹VOCdevkit\VOC2007，里面有以下四项

├── Annotations（标注的xml文件）
       ├── object_00001.xml
       └── object_00002.xml

├──ImageSets

        ├──Main

                ├──train.txt（自己划分出来的训练集图片名称，没有jpg后缀的）

                ├──val.txt（列举验证集图片的名称）

                ├──test.txt

├── JPEGImages（训练图片）
        ├── object_00001.jpg
        └── object_00002.jpg
         …
├── Labels（这个里面的内容后面voc_label.py生成）
       ├── object_00001.txt
       └── object_00002.txt
        …
然后把VOCdevkit存入darknet-master/build/darknet/x64/data中

然后修改同文件夹下的voc/voc_label.py。（我是在x64下运行，当前文件夹下会生成2007_train.txt，2007_val.txt，2007_test.txt）

我这里的代码修改为

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

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
#项目类别
classes = ["stone", "fod"]


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):
#这里改为.xml文件夹的路径
    in_file = open('data/VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
#这里是生成每张图片对应的.txt文件的路径 
    out_file = open('data/VOCdevkit/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('data/VOCdevkit/VOC%s/labels/'%(year)):
        os.makedirs('data/VOCdevkit/VOC%s/labels/'%(year))
    image_ids = open('data/VOCdevkit/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/data/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
        convert_annotation(year, image_id)
    list_file.close()

然后cd到darknet-master下的cfg中，复制coco.data到本地，更改为obj.data，将其更改为：

classes= 2
#之前生成两个文件的地址，相对地址也可
train  = /home/xinwen/workspace/darknet-master/2007_train.txt
valid  = /home/xinwen/workspace/darknet-master/2007_val.txt
names = cfg/obj.names
backup = backup/

coco.names也同样复制到当前文件夹下，命名为obj.names，改为：

#类别名称
stone
fod

接下来是yolov4.custom.cfg复制到当前，命名为yolo-obj.cfg(等下的训练命令以此)，改动以下： 

 修改classes和上面的filters，filters=（类别+5）*3，一共有三个要修改地方

然后还有max_batches改为6000（类别比较少，不用那么久）

steps改为4800,5400（一般分别为batches的80%和90%）

还有width和hight改为416

第6行，第7行的batch和subdivisions可以视服务器的性能而定，这里改为

batch=16

subdivisions=32

然后下载预训练集yolov4.conv.137到darknet-master

然后就可以开始训练了，在darknet-master下

./darknet detector train cfg/obj.data cfg/yolo-obj.cfg yolov4.conv.137

测试

./darknet detector test cfg/obj.data cfg/yolo-obj.cfg backup/yolo-obj_last.weights