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