YOLOV3目標檢測模型訓練實例

YOLOV3目標檢測

從零開始學習使用keras-yolov3進行圖片的目標檢測，比較詳細地記錄了準備以及訓練過程，提供一個信號燈的目標檢測模型訓練實例，並提供相關代碼與訓練集。

DEMO測試

YOLO提供了模型以及源碼，首先使用YOLO訓練好的權重文件進行快速測試，首先下載權重文件

https://pjreddie.com/media/files/yolov3.weights

將yolo3的版本庫clone到本地，本次測試的commit id爲e6598d1

git clone [email protected]:qqwweee/keras-yolo3.git

安裝各種依賴，缺啥就安啥，注意依賴版本對應，以下版本僅供參考

Keras==2.2.4
numpy==1.16.0
tensorflow==1.12.0
...

執行convert.py文件，將darknet的yolo轉換爲可以用於keras的h5文件，生成的文件被保存在model_data下，此外convert.py和yolov3.vfg在git clone後的根目錄已經給出，不需要單獨下載。

python convert.py yolov3.cfg yolov3.weights model_data/yolo.h5

使用python yolo_video.py -h獲取help內容

usage: yolo_video.py [-h] [--model MODEL] [--anchors ANCHORS]
                     [--classes CLASSES] [--gpu_num GPU_NUM] [--image]
                     [--input [INPUT]] [--output [OUTPUT]]
optional arguments:
  -h, --help         show this help message and exit
  --model MODEL      path to model weight file, default model_data/yolo.h5
  --anchors ANCHORS  path to anchor definitions, default
                     model_data/yolo_anchors.txt
  --classes CLASSES  path to class definitions, default
                     model_data/coco_classes.txt
  --gpu_num GPU_NUM  Number of GPU to use, default 1
  --image            Image detection mode, will ignore all positional
                     arguments
  --input [INPUT]    Video input path
  --output [OUTPUT]  [Optional] Video output path

本次測試是進行圖片的目標檢測，注意當參數爲--image時會忽略所有位置參數，也就是說當進行圖片檢測時每次都需要手動輸入位置，當然這可以以後通過自行構建代碼修改

python yolo_video.py --image

之後會出現Input image filename:我是放到./img/3.jpg下，於是就直接將路徑輸入

稍等一會就可以識別完成

模型訓練

準備數據集

首先需要準備好目錄結構，可以在 http://host.robots.ox.ac.uk/pascal/VOC/voc2007/ 中下載VOC2007數據集，然後刪除其中所有的文件，僅保留目錄結構，也可以手動建立如下目錄結構

然後將所有的圖片放置在JPEGImages目錄下，然後在
https://github.com/tzutalin/labelImg 下載labelImg標註工具，此工具是爲了將圖片框選標註後生成XML文件，使用labelImg打開圖片，標註好後將圖片生成的XML文件放置於Annotations文件夾內，保存的名字就是圖片的名字。

準備訓練文件

在VOCdevkit/VOC2007下建立一個python文件，將代碼寫入並運行，即會在VOCdevkit/VOC2007/ImageSets/Main下生成四個txt文件

import os
import random
 
trainval_percent = 0
train_percent = 1 # 全部劃分爲訓練集，因爲yolo3在訓練時依舊會劃分訓練集與測試集，不需要在此劃分
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets/Main'
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)
 
ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('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:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)
 
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

在VOCdevkit的上層目錄，我目前的目錄結構爲Train下，建立python文件並運行，生成三個txt文件，注意，此處代碼需要將classes更改成需要訓練的類別，我只需要訓練person一類，所以此處數組中只有person類別

import xml.etree.ElementTree as ET
from os import getcwd

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]

classes = ["person"]


def convert_annotation(year, image_id, list_file):
    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id),'rb')
    tree=ET.parse(in_file)
    root = tree.getroot()

    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 = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text))
        list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))

wd = getcwd()

for year, image_set in sets:
    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s.txt'%(image_set), 'w')
    for image_id in image_ids:
        list_file.write('VOCdevkit/VOC%s/JPEGImages/%s.jpg'%(year, image_id))
        convert_annotation(year, image_id, list_file)
        list_file.write('\n')
    list_file.close()

接下來將Train目錄下所有的文件複製到git clone後的目錄下，此時的文件目錄結構是這樣的

修改參數

此時需要修改model_data/coco_classes.txt與voc_classes.txt文件，這兩個文件都是需要存放訓練類別的，同樣我只是訓練person類別，此處只有一行person。
接下來修改yolov3.cfg，假如你不需要加載預訓練的權重，那麼此文件是沒有必要修改的，此文件是爲生成yolo_weights.h5作配置的，在此文件中搜索yolo，會有三處匹配，都是相同的更改方式，以第一次匹配舉例，三處註釋位置，也就是共需改動9行

...
[convolutional]
size=1
stride=1
pad=1
filters=18  # 3*(5+len(classes)) # 我訓練一種類別 即 3*(5+1) = 18
activation=linear


[yolo]
mask = 6,7,8
anchors = 10,13,  16,30,  33,23,  30,61,  62,45,  59,119,  116,90,  156,198,  373,326
classes=1 # 一種類別
num=9
jitter=.3
ignore_thresh = .5
truth_thresh = 1
random=1 # 顯存小就改爲0   
...

運行python convert.py -w yolov3.cfg yolov3.weights model_data/yolo_weights.h5生成model_data / yolo_weights.h5用於加載預訓練的權重。

訓練模型

之後就可以開始訓練了，因爲我一開始暫時沒有數據，就隨便找了幾張圖片標註後試了一下，因爲不足十張，外加我在構建VOC數據集時又劃分了一下數據集與訓練集，而train.py又默認將數據劃分了0.1的訓練集，不足十張乘0.1取整就是0，導致我一直報錯，此處一定要注意，一定要有驗證集，也就是至少需要有兩張圖片，一張作爲訓練集一張作爲驗證集，否則運行train.py時會報錯KeyError: 'val_loss'，運行train_bottleneck.py會報錯IndexError: list index out of range，此外還需要注意的是需要手動建立logs/000/目錄，防止保存模型時無法找到目錄而拋出異常。訓練一般使用train.py就可以了，對於出現的問題多多去看看github的issue與README，很多問題都會有討論與解決，對於train.py我略微做了一些更改以適應我的訓練目的，對於一些更改的地方有註釋

"""
Retrain the YOLO model for your own dataset.
"""

import numpy as np
import keras.backend as K
from keras.layers import Input, Lambda
from keras.models import Model
from keras.optimizers import Adam
from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping

from yolo3.model import preprocess_true_boxes, yolo_body, tiny_yolo_body, yolo_loss
from yolo3.utils import get_random_data


def _main():
    annotation_path = 'train.txt'
    log_dir = 'logs/000/'
    classes_path = 'model_data/voc_classes.txt'
    anchors_path = 'model_data/yolo_anchors.txt'
    class_names = get_classes(classes_path)
    num_classes = len(class_names)
    anchors = get_anchors(anchors_path)

    input_shape = (416,416) # multiple of 32, hw

    # 此處去掉了 create_tiny_model 的判斷 # load_pretrained 爲False即不加載預訓練的權重，爲True則加載預訓練的權重
    model = create_model(input_shape, anchors, num_classes,load_pretrained=False,
            freeze_body=2, weights_path='model_data/yolo_weights.h5') # make sure you know what you freeze

    logging = TensorBoard(log_dir=log_dir)

    # ModelCheckpoint 回調檢查模型週期 更改爲每10次檢查
    checkpoint = ModelCheckpoint(log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
        monitor='val_loss', save_weights_only=True, save_best_only=True, period=10)
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1)
    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=6000, verbose=1)

    # 對輸入劃分訓練集與測試集的比重
    val_split = 0.3
    with open(annotation_path) as f:
        lines = f.readlines()
    np.random.seed(10101)
    np.random.shuffle(lines)
    np.random.seed(None)
    num_val = int(len(lines)*val_split)
    num_train = len(lines) - num_val

    # Train with frozen layers first, to get a stable loss.
    # Adjust num epochs to your dataset. This step is enough to obtain a not bad model.
    if True:
        model.compile(optimizer=Adam(lr=1e-3), loss={
            # use custom yolo_loss Lambda layer.
            'yolo_loss': lambda y_true, y_pred: y_pred})

        # batch_size 需要針對顯存更改數量
        batch_size = 10
        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))

        # epochs 即訓練次數
        model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
                steps_per_epoch=max(1, num_train//batch_size),
                validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
                validation_steps=max(1, num_val//batch_size),
                epochs=50, 
                initial_epoch=0,
                callbacks=[logging, checkpoint])
        model.save_weights(log_dir + 'trained_weights_stage_1.h5')

    # Unfreeze and continue training, to fine-tune.
    # Train longer if the result is not good.
    if True:
        for i in range(len(model.layers)):
            model.layers[i].trainable = True
        model.compile(optimizer=Adam(lr=1e-4), loss={'yolo_loss': lambda y_true, y_pred: y_pred}) # recompile to apply the change
        print('Unfreeze all of the layers.')

        # batch_size 需要針對顯存更改數量
        batch_size = 10 # note that more GPU memory is required after unfreezing the body
        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))

        # epochs即訓練次數
        model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
            steps_per_epoch=max(1, num_train//batch_size),
            validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
            validation_steps=max(1, num_val//batch_size),
            epochs=50,
            initial_epoch=50)
        model.save_weights(log_dir + 'trained_weights_final.h5')

    # Further training if needed.


def get_classes(classes_path):
    '''loads the classes'''
    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names

def get_anchors(anchors_path):
    '''loads the anchors from a file'''
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = [float(x) for x in anchors.split(',')]
    return np.array(anchors).reshape(-1, 2)


def create_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
            weights_path='model_data/yolo_weights.h5'):
    '''create the training model'''
    K.clear_session() # get a new session
    image_input = Input(shape=(None, None, 3))
    h, w = input_shape
    num_anchors = len(anchors)

    y_true = [Input(shape=(h//{0:32, 1:16, 2:8}[l], w//{0:32, 1:16, 2:8}[l], \
        num_anchors//3, num_classes+5)) for l in range(3)]

    model_body = yolo_body(image_input, num_anchors//3, num_classes)
    print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))

    if load_pretrained:
        model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
        print('Load weights {}.'.format(weights_path))
        if freeze_body in [1, 2]:
            # Freeze darknet53 body or freeze all but 3 output layers.
            num = (185, len(model_body.layers)-3)[freeze_body-1]
            for i in range(num): model_body.layers[i].trainable = False
            print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))

    model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
        arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})(
        [*model_body.output, *y_true])
    model = Model([model_body.input, *y_true], model_loss)

    return model

def create_tiny_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
            weights_path='model_data/tiny_yolo_weights.h5'):
    '''create the training model, for Tiny YOLOv3'''
    K.clear_session() # get a new session
    image_input = Input(shape=(None, None, 3))
    h, w = input_shape
    num_anchors = len(anchors)

    y_true = [Input(shape=(h//{0:32, 1:16}[l], w//{0:32, 1:16}[l], \
        num_anchors//2, num_classes+5)) for l in range(2)]

    model_body = tiny_yolo_body(image_input, num_anchors//2, num_classes)
    print('Create Tiny YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))

    if load_pretrained:
        model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
        print('Load weights {}.'.format(weights_path))
        if freeze_body in [1, 2]:
            # Freeze the darknet body or freeze all but 2 output layers.
            num = (20, len(model_body.layers)-2)[freeze_body-1]
            for i in range(num): model_body.layers[i].trainable = False
            print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))

    model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
        arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})(
        [*model_body.output, *y_true])
    model = Model([model_body.input, *y_true], model_loss)

    return model

def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes):
    '''data generator for fit_generator'''
    n = len(annotation_lines)
    i = 0
    while True:
        image_data = []
        box_data = []
        for b in range(batch_size):
            if i==0:
                np.random.shuffle(annotation_lines)
            image, box = get_random_data(annotation_lines[i], input_shape, random=True)
            image_data.append(image)
            box_data.append(box)
            i = (i+1) % n
        image_data = np.array(image_data)
        box_data = np.array(box_data)
        y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
        yield [image_data, *y_true], np.zeros(batch_size)

def data_generator_wrapper(annotation_lines, batch_size, input_shape, anchors, num_classes):
    n = len(annotation_lines)
    if n==0 or batch_size<=0: return None
    return data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes)

if __name__ == '__main__':
    _main()

測試模型

當模型訓練完成後，就可以加載模型進行圖片測試了

import sys
import argparse
from yolo import YOLO, detect_video
from PIL import Image


if __name__ == '__main__':
    config = {
        "model_path": "logs/000/trained_weights_final.h5", # 加載模型
        "score": 0.1, # 超出這個值的預測纔會被顯示
        "iou": 0.5, # 交併比
    }
    yolo = YOLO(**config)
    image = Image.open("./img/1.jpg")
    r_image = yolo.detect_image(image)
    r_image.save("./img/2.jpg")

此後就需要不斷開始優化參數並訓練了，其實在目錄中有很多文件是用不到的或者是使用一次後就一般不會再用到了，可以備份一下代碼後適當精簡目錄結構。

模型訓練實例

從百度下載了50張信號燈的圖片作訓練集，實例僅爲模型訓練的Demo，數據集比較小，相關信息僅供參考。

運行環境

cuda 8.0
python 3.6
keras 2.1.5
tensorflow-gpu 1.4.0

相關配置

val_split = 0.1 # 訓練集與測試集劃分比例
batch_size = 5 # 每次訓練選擇樣本數
epochs = 300 # 訓練三百次

運行結果

數據集中的紅燈比較多，所以訓練結果中紅燈的置信度爲0.60和0.72，綠燈樣本較少，識別的綠燈的置信度爲0.38，整體效果還算可以。

loss: 25.8876 - val_loss: 38.1282

原圖

識別

實例代碼

如果覺得不錯，點個star吧 😃
https://github.com/WindrunnerMax/Yolov3-Train