目標檢測 YOLO v3 訓練 人臉檢測模型

SpikeKing發表於2018-07-05
YOLO模型

YOLO，是You Only Look Once的縮寫，一種基於深度卷積神經網路的物體檢測演算法，YOLO v3是YOLO的第3個版本，檢測演算法更快更準。

本文原始碼：https://github.com/SpikeKing/keras-yolo3-detection

歡迎Follow我的GitHub：https://github.com/SpikeKing

YOLO v3已經提供 COCO（Common Objects in Context）資料集的模型引數。我們可以把COCO的模型引數作為預訓練引數，再結合已有的資料集，建立自己的檢測演算法。

本例使用WIDER FACE人臉資料，訓練一個高精度的人臉檢測模型。

WIDER

資料集：WIDER Face

WIDER

建立時間：2015-11-19

WIDER FACE 資料集是一個人臉檢測基準（benchmark）資料集，圖片選取自 WIDER（Web Image Dataset for Event Recognition） 資料集。圖片數 32,203 張，人臉數 393,703 個，在大小（scale）、位置（pose）、遮擋（occlusion）等不同形式中，人臉是高度變換的。WIDER FACE 資料集是基於61個事件類別，每個事件類別，隨機選取訓練40%、驗證10%、測試50%。訓練和測試含有邊框（bounding box）真值（ground truth），而驗證不含。

資料集在官網可以公開下載，其中在Face annotations 中，wider_face_train_bbx_gt.txt是邊框真值，資料格式如下：

0--Parade/0_Parade_marchingband_1_849.jpg
1
449 330 122 149 0 0 0 0 0 0 
複製程式碼

資料說明：

  • 第1行：圖片的位置和名稱；
  • 第2行：邊框的數量；
  • 第3~n行：每個人臉的邊框和屬性：
    • 其中1~4位是x1, y1, w, h
    • blur：模糊，0清晰、1一般、2嚴重；
    • expression：表情，0正常、1誇張；
    • illumination：曝光，0正常、1極度；
    • occlusion：遮擋，0無、1部分、2大量；
    • pose：姿勢，0正常，1非典型；

wider_face_val_bbx_gt.txt與此類似。

圖片資料的清晰度一般，大小不一，尺寸為1024x，寬度相同。

資料轉換

為了符合訓練要求，需要轉換wider資料集中的邊框格式，為訓練要求的邊框格式。

即檔案路徑，邊框xmin,ymin,xmax,ymax,label：

data/WIDER_val/images/10--People_Marching/10_People_Marching_People_Marching_2_433.jpg 614,346,771,568,0 245,382,392,570,0 353,222,461,390,0 498,237,630,399,0
複製程式碼

轉換原始碼。遍歷資料資料夾，逐行解析不同格式的資料，寫入檔案。注意：

  1. 物體框，Wider的資料格式是x,y,w,h，而訓練的資料格式是xmin,ymin,xmax,ymax；
  2. 只檢測人臉一個類別，則類別索引均為0；

具體參考工程的wider_annotation.py指令碼。

def generate_train_file(bbx_file, data_folder, out_file):
    paths_list, names_list = traverse_dir_files(data_folder)
    name_dict = dict()
    for path, name in zip(paths_list, names_list):
        name_dict[name] = path

    data_lines = read_file(bbx_file)

    sub_count = 0
    item_count = 0
    out_list = []

    for data_line in data_lines:
        item_count += 1
        if item_count % 1000 == 0:
            print('item_count: ' + str(item_count))

        data_line = data_line.strip()
        l_names = data_line.split('/')
        if len(l_names) == 2:
            if out_list:
                out_line = ' '.join(out_list)
                write_line(out_file, out_line)
                out_list = []

            name = l_names[-1]
            img_path = name_dict[name]
            sub_count = 1
            out_list.append(img_path)
            continue

        if sub_count == 1:
            sub_count += 1
            continue

        if sub_count >= 2:
            n_list = data_line.split(' ')
            x_min = n_list[0]
            y_min = n_list[1]
            x_max = str(int(n_list[0]) + int(n_list[2]))
            y_max = str(int(n_list[1]) + int(n_list[3]))
            p_list = ','.join([x_min, y_min, x_max, y_max, '0'])  # 標籤全部是0，人臉
            out_list.append(p_list)
            continue
複製程式碼

類別檔案wider_classes.txt僅有一行，就是face。

訓練

YOLO v3的訓練過程，引數：標註資料、類別、儲存路徑、預訓練模型、anchors、輸入尺寸。

annotation_path = 'dataset/WIDER_train.txt'  # 資料
classes_path = 'configs/wider_classes.txt'  # 類別

log_dir = 'logs/002/'  # 日誌資料夾

pretrained_path = 'model_data/yolo_weights.h5'  # 預訓練模型
anchors_path = 'configs/yolo_anchors.txt'  # anchors

class_names = get_classes(classes_path)  # 類別列表
num_classes = len(class_names)  # 類別數
anchors = get_anchors(anchors_path)  # anchors列表

input_shape = (416, 416)  # 32的倍數，輸入影象
複製程式碼

建立模型：

  1. input_shape是輸入影象的尺寸；
  2. anchors是檢測框的尺寸；
  3. num_classes是類別數；
  4. freeze_body，模式1是全部凍結，模式2是訓練最後三層；
  5. weights_path，預訓練權重的路徑；
  6. logging是TensorBoard的回撥，checkpoint是儲存權重的回撥；
model = create_model(input_shape, anchors, num_classes,
                     freeze_body=2,
                     weights_path=pretrained_path)  # make sure you know what you freeze

logging = TensorBoard(log_dir=log_dir)
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=3)  # 只儲存weights，
複製程式碼

訓練資料與驗證資料：

val_split = 0.1  # 訓練和驗證的比例
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  # 訓練集數量
複製程式碼

模型編譯與fit資料：

  1. 損失函式只使用y_pred預測結果；
  2. 批次數是32個；
  3. 訓練資料與驗證資料都來源於data_generator_wrapper；
  4. 訓練過程中，通過checkpoint儲存權重，最終再儲存最終的權重；
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 = 32  # batch尺寸
print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
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=200,
                    initial_epoch=0,
                    callbacks=[logging, checkpoint])
model.save_weights(log_dir + 'trained_weights_stage_1.h5')  # 儲存最終的引數，再訓練過程中，通過回撥儲存
複製程式碼

模型建立：

  1. 建立YOLO v3的模型，yolo_body，引數影象輸入、每個尺度的anchor數、類別數；
  2. 載入預訓練權重，凍結引數，保留最後三層；
  3. 自定義Lambda，模型的損失函式層；
  4. 輸入是YOLO模型輸入和真值，輸出是損失函式；
def create_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
                 weights_path='model_data/yolo_weights.h5'):
    K.clear_session()  # 清除session
    image_input = Input(shape=(None, None, 3))  # 圖片輸入格式
    h, w = input_shape  # 尺寸
    num_anchors = len(anchors)  # anchor數量

    # YOLO的三種尺度，每個尺度的anchor數，類別數+邊框4個+置信度1
    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)  # model
    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)  # 模型，inputs和outputs

    return model
複製程式碼

資料生成器：

  1. data_generator_wrapper用於條件檢查；
  2. 將輸入的標註行隨機；
  3. 根據批次數，將影象放入image_data中，將邊框和引數放入真值y_true中；
  4. 輸出影象和邊框，還有批次數的填充，用於儲存置信度。
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)
複製程式碼

具體原始碼參考yolo3_train.py，即可生成人臉檢測模型。

驗證

在yolo3_predict.py中，替換已訓練好的模型引數：

模型：ep108-loss44.018-val_loss43.270.h5

檢測圖片：

output

其他：增加更多的資料集，與具體的需求圖片結合，都可以提升檢測效果。

OK, that's all! Enjoy it!

相關文章