YOLOv3训练自己的数据

 

1.  下载预训练权重文件

YOLOv3使用在Imagenet上预训练好的模型参数（文件名称： darknet53.conv.74，大小76MB）基础上继续训练。
darknet53.conv.74下载链接：  https://pjreddie.com/media/files/darknet53.conv.74，下载完成后放在darknet主目录。

也可以直接在darknet目录下通过wget命令下载：

wget https://pjreddie.com/media/files/darknet53.conv.74

2.  准备打标工具并对自己的图片数据打标

打标工具推荐使用 labelImg，下载地址：https://github.com/tzutalin/labelImg 或 http://download.csdn.net/download/dcrmg/9974195
labelImg使用很简单，在图片的物体上画框然后给一个标签就可以了，打标结果的保存格式是xml文件。
例如对于train1.jpg，打标结果保存为train1.xml

 

3.  xml标签文件格式转换

YOLO训练的标签文件是txt格式，需要把第2步中的xml文件转换。

•  1） 在darknet主目录下创建4个文件夹： trainImage、validateImage、trainImageXML 和 validateImageXML，并分别存放第2步中的训练集图片、验证集图片、训练集xml标签和验证集xml标签
•  2） 借助createID.py生成训练集和验证集的图片名称列表trainImageId.txt和validateImageId.txt。

     createID.py 代码：

# -*- coding: utf-8 -*-
import os;
import shutil;
 
def listname(path,idtxtpath):
    filelist = os.listdir(path);  # 该文件夹下所有的文件（包括文件夹）
    filelist.sort()
    f = open(idtxtpath, 'w');
    for files in filelist:  # 遍历所有文件
        Olddir = os.path.join(path, files);  # 原来的文件路径
        if os.path.isdir(Olddir):  # 如果是文件夹则跳过
            continue;
        f.write(files);
        f.write('\n');
    f.close();
 
savepath = os.getcwd()
imgidtxttrainpath = savepath+"/trainImageId.txt"
imgidtxtvalpath = savepath + "/validateImageId.txt"
listname(savepath + "/trainImage",imgidtxttrainpath)
listname(savepath + "/validateImage",imgidtxtvalpath)
print "trainImageId.txt && validateImageId.txt have been created!"

  3） 借助trans.py生成训练集和验证集的完整路径列表并完成标签xml文件到txt文件的转换
  trans.py代码：

import xml.etree.ElementTree as ET
import pickle
import string
import os
import shutil
from os import listdir, getcwd
from os.path import join
import cv2
 
sets=[('2012', 'train')]
 
classes = ["class1","class2","class3","class4"]
 
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(image_id,flag,savepath):
 
    if flag == 0:
        in_file = open(savepath+'/trainImageXML/%s.xml' % (os.path.splitext(image_id)[0]))
        out_file = open(savepath+'/trainImage/%s.txt' % (os.path.splitext(image_id)[0]), 'w')
        tree = ET.parse(in_file)
        root = tree.getroot()
        size = root.find('size')
 
        img = cv2.imread('./trainImage/'+str(image_id))
        h = img.shape[0]
        w = img.shape[1]
 
    elif flag == 1:
        in_file = open(savepath+'/validateImageXML/%s.xml' % (os.path.splitext(image_id)[0]))
        out_file = open(savepath+'/validateImage/%s.txt' % (os.path.splitext(image_id)[0]), 'w')
 
        tree = ET.parse(in_file)
        root = tree.getroot()
        size = root.find('size')
 
        img = cv2.imread('./validateImage/' + str(image_id))
        h = img.shape[0]
        w = img.shape[1]
 
    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:
    savepath = os.getcwd();
    idtxt = savepath + "/validateImageId.txt";
    pathtxt = savepath + "/validateImagePath.txt";
    image_ids = open(idtxt).read().strip().split()
    list_file = open(pathtxt, 'w')
    s = '\xef\xbb\xbf'
    for image_id in image_ids:
        nPos = image_id.find(s)
        if nPos >= 0:
            image_id = image_id[3:]
        list_file.write('%s/validateImage/%s\n' % (wd, image_id))
        print(image_id)
        convert_annotation(image_id, 1, savepath)
    list_file.close()
 
    idtxt = savepath + "/trainImageId.txt";
    pathtxt = savepath + "/trainImagePath.txt" ;
    image_ids = open(idtxt).read().strip().split()
    list_file = open(pathtxt, 'w')
    s = '\xef\xbb\xbf'
    for image_id in image_ids:
        nPos = image_id.find(s)
        if nPos >= 0:
           image_id = image_id[3:]
        list_file.write('%s/trainImage/%s\n'%(wd,image_id))
        print(image_id)
        convert_annotation(image_id,0,savepath)
    list_file.close()

注意： 需要根据自己的类别更改trans.py文件第12行的classes，有几个类别写几个。
执行之后在darknet主目录下生成trainImagePath.txt、validateImagePath.txt和所有的txt标注文件。

 

4. 修改配置文件

•   1） 修改data/voc.names 文件

     把 voc.names文件内容改成自己的分类，例如有3个分类class_1,class_2,class_3,则voc.names内容改为：
      class_1
      class_2
      class_3

 

•   2) 修改cfg/voc.data文件

  根据自己的实际情况做以下修改：
   classes = N       #（N为自己的分类数量，如有10类不同的对象，N = 10）
   train = /home/XXX/darknet/trainImagePath.txt    # 训练集完整路径列表
   valid = /home/XXX/darknet/validateImagePath.txt   # 测试集完整路径列表
   names = data/voc.names    # 类别文件
   backup = backup     #（训练结果保存在darknet/backup/目录下）

 

•    3） 修改cfg/yolov3-voc.cfg 文件

   1. classes = N （N为自己的分类数）
   2. 修改每一个[yolo]层（一共有3处）之前的filters为 3*(classes+1+4)，如有3个分类，则修改 filters = 24
   3. (可选) 修改训练的最大迭代次数， max_batches = N

 

5. YOLOv3训练

./darknet detector train cfg/voc.data cfg/yolov3-voc.cfg darknet53.conv.74

训练完成后结果文件 ‘yolov3-voc_final.weights’ 保存在 backup文件中。

 

6. 自训练模型测试

./darknet detector test cfg/voc.data cfg/yolov3-voc.cfg backup/yolov3-voc_final.weights 01.jpg