YOLOv4-tiny 模型訓練教學

tags: YOLO

本篇記錄如何使用自己的資料集，利用YOLO進行訓練

Step 0: Environment setting

1. 下載 darknet

git clone https://github.com/AlexeyAB/darknet.git

2. 編譯 Yolo (根據自身電腦環境，修改 Makefile)
   參考 https://github.com/AlexeyAB/darknet#how-to-compile-on-linux-using-make

Step 1: Label data

安裝 labelImg (Windows 作業系統)
https://drive.google.com/drive/folders/13f6XFDPI-cguUk9UUH9SH5AqrDU1gkTp?usp=share_link

開啟 labelImg 軟體後:
Step 1. 選擇資料夾

選取欲訓練的照片資料夾，並選取儲存 Label 的資料夾。(儲存Label 的資料夾，建議與訓練的照片資料夾相同)

Step 2. 選擇儲存成 Yolo 格式

Step 3. 標記目標物件
點選 Create RectBox，框選目標物件，並輸入物件名稱

Step 4. 存檔

此時，會根據 xxx.jpg ，產生對應的 xxx.txt。
xxx.txt 為Yolo格式，作為訓練 Yolo 模型的檔案。

完成所有資料的標記後，將所有訓練的照片檔案(.jpg)以及標記檔案(.txt)，放至同一份資料夾。

Step 2: Data preprocessing (Option)

此步驟為在 Step 1. Label data，使用 VOC 作為 label 的儲存格式(.xml)，而非 Yolo 格式(.txt)，則需要將 xml 檔案轉 txt　檔案，以符合　Yolo 的訓練格式。
若 Step 1. Label data 有正確將 label 設為 Yolo 格式(.txt)，則不必進行此步驟。

定義getYoloFormat函數，功用如下

1. 把物件類別轉成數字
2. 把 images 裡的照片重新命名，並存入名為 "yolo" 的資料夾
3. 把 labels 裡的 xml 轉成 txt 並重新命名存入名為 "yolo" 的資料夾
4. xxx.jpg 和 xxx.tx 名稱會一一對應

程式碼如下，並將status_dic內改成對應的label

from bs4 import BeautifulSoup
import os
import shutil
from IPython.display import clear_output

status_dic = {'person':0}             #用dictionary 記錄label的名稱

def getYoloFormat(filename, label_path, img_path, yolo_path, newname):
    with open(label_path+filename, 'r') as f:
        soup = BeautifulSoup(f.read(), 'xml')
        imgname = soup.select_one('filename').text         #讀取xml
        image_w = soup.select_one('width').text
        image_h = soup.select_one('height').text
        ary = []
        for obj in soup.select('object'):                  #取出xmin, xmax, ymin, ymax及name
            xmin = int(obj.select_one('xmin').text)          #並且用status_dictionary 來轉換name，good =>2
            xmax = int(obj.select_one('xmax').text)
            ymin = int(obj.select_one('ymin').text)
            ymax = int(obj.select_one('ymax').text)
            objclass = status_dic.get(obj.select_one('name').text)

            x = (xmin + (xmax-xmin)/2) * 1.0 / float(image_w)    #YOLO吃的參數檔有固定的格式
            y = (ymin + (ymax-ymin)/2) * 1.0 / float(image_h)    #先照YOLO的格式訂好x,y,w,h
            w = (xmax-xmin) * 1.0 / float(image_w)
            h = (ymax-ymin) * 1.0 / float(image_h)
            ary.append(' '.join([str(objclass), str(x), str(y), str(w), str(h)]))
      
        if os.path.exists(img_path+imgname+'.jpg'):                              # 圖片本來在image裡面，把圖片移到yolo資料夾下    
            shutil.copyfile(img_path+imgname+'.jpg', yolo_path+newname+'.jpg')     #同時把yolo參數檔寫到yolo之下
            with open(yolo_path+newname+'.txt', 'w') as f:
                f.write('\n'.join(ary))            
        elif os.path.exists(img_path+imgname):  #有的labelImg名稱已經自動加上.jpg                          
            shutil.copyfile(img_path+imgname, yolo_path+newname+'.jpg')
            with open(yolo_path+newname+'.txt', 'w') as f:
                f.write('\n'.join(ary))
                
def update_progress(progress):
    bar_length = 20
    if isinstance(progress, int):
        progress = float(progress)
    if not isinstance(progress, float):
        progress = 0
    if progress < 0:
        progress = 0
    if progress >= 1:
        progress = 1
    block = int(round(bar_length * progress))
    clear_output(wait = True)
    text = "Progress: [{0}] {1:.1f}%".format( "#" * block + "-" * (bar_length - block), progress * 100)
    print(text)    

呼叫getYoloFormat

labelpath = 'darknet/labels'  # label xml資料夾         
imgpath = 'darknet/images' # image 資料夾
yolopath = 'darknet/yolo' # 存放轉換成yolo訓練格式的照片及txt檔案

total_progress = len(os.listdir(labelpath))
progress = 0

for idx, f in enumerate(os.listdir(labelpath)):   #透過getYoloFormat將圖像和參數檔全部寫到YOLO下
    progress += 1
    try:
        if f.split('.')[1] == 'xml':
            getYoloFormat(f, labelpath, imgpath, yolopath, str(idx))
        else:
            print('Error, the file should be the xml')
    except Exception as e:
        print(e)
    
    update_progress(progress/total_progress)

Step 3: 切割訓練及測試資料集

範例程式碼

import os
import random

# 存放 yolo 要訓練的照片及txt檔案的資料夾 
datasets = ['絕對路徑/yolo/'+ f for f in os.listdir('./yolo/') if  f.endswith('.jpg')]
random.shuffle(datasets)


len_dataset = int(len(datasets) * 0.7)
random.shuffle(datasets)
print('All data length',len(datasets))
print(len_dataset) 
print(len(datasets)-len_dataset)     

with open('/root/notebooks/yolov4/darknet/cfg/train.txt', 'w') as f:
    f.write('\n'.join(datasets[:len_dataset]))
    
with open('/root/notebooks/yolov4/darknet/cfg/test.txt', 'w') as f:
    f.write('\n'.join(datasets[len_dataset:]))

Step 4: 編輯訓練參數檔案

此步驟將建立以下三個檔案:

• obj.data

建立obj.data檔案

說明：
classes：總共的物件數量
train: 包含所有欲訓練的照片及標記檔案路徑
valid: 包含所有欲測試的照片及標記檔案路徑
names: obj.names 路徑位置
backup: 存放模型的訓練權重的資料夾位置

範例：

classes= 1
train  = cfg/train.txt
valid  = cfg/test.txt
names  = cfg/obj.names
backup = backup 

• obj.name

建立obj.name檔案，輸入所有的label種類。
注意: label 名稱須按照 Step 1. 的標記順序擺放。
例如：

person

• obj.cfg

以 Yolov4-tiny 為範例，./darknet/cfg/ 底下，修改 yolov4-tiny-custom.cfg，複製一份新檔案，重新命名為 obj.cfg，並修改 filter 以及 classes :

修改方式:
:bulb: 尋找關鍵字:yolo，並參考下圖的相對位置做更改
:bulb: 修改規則為 [yolo] 內的classes以及 [yolo] 標籤前的filters，如下圖

filter數量計算方式為: 3*(classes + 5)

參考資料: Yolo 官方網站
https://github.com/AlexeyAB/darknet#how-to-train-to-detect-your-custom-objects

Step 5: Train

• 下載 pre-trained model (以 Yolov4-tiny 為範例)
  https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v4_pre/yolov4-tiny.conv.29

• 開始訓練
  於 darknet 底下執行 (需改成對應的參數路徑)

./darknet detector train cfg/obj.data cfg/obj.cfg yolov4-tiny.conv.29 -map -dont_show

Step 6: Test

查看辨識結果，產生於darknet/prediction.jpg

./darknet detector test cfg/obj.data cfg/obj.cfg cfg/your_weight.weights

(option) 設置 thresh ，可查看更低機率的class

./darknet detector test cfg/obj.data cfg/obj.cfg cfg/your_weight.weights -thresh 0.1

:bulb:TroubleShooting

辨識不出東西，設 -thresh 0 ，發現預測機率都是 0% 時，將darknet/MakeFile中的cudnn設為0，並重新make clean 以及 make

番外 : 儲存log檔，並查看訓練過程的loss以及IOU

• 於darknet底下建立visulization資料夾，執行以下指令開始訓練

./darknet detector train cfg/obj.data cfg/yolov3_PB.cfg cfg/darknet53.conv.74 2>1 | tee visualization/yolov3_PB.log

• 讀取檔案

import inspect
import os
import random
import sys
def extract_log(log_file, new_log_file, key_word):
    with open(log_file, 'r') as f:
        with open(new_log_file, 'w') as train_log:
            for line in f:
                # 去除多gpu的同步log
                if 'Syncing' in line:
                    continue
                # 去除除零错误的log
                if 'nan' in line:
                    continue
                if key_word in line:
                    train_log.write(line)
    f.close()
    train_log.close()

extract_log('yolov3_PH.log','train_log_loss.txt','images')
extract_log('yolov3_PH.log','train_log_iou.txt','IOU')

• avg loss

#!/usr/bin/python
#coding=utf-8

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 根据log_loss.txt的行数修改lines, 修改训练时的迭代起始次数(start_ite)和结束次数(end_ite)。
lines = 3000
start_ite = 0  # train_log_loss.txt里面的最小迭代次数
end_ite = 3000  # train_log_loss.txt里面的最大迭代次数
step = 10  # 跳行数，决定画图的稠密程度
igore = 100  # 当开始的loss较大时，忽略前igore次迭代，注意这里是迭代次数

y_ticks = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0, 2.1, 2.2, 2.3]  # 纵坐标的值
data_path =  'train_log_loss.txt' # log_loss的路径。
result_path = 'avg_loss' # 保存结果的路径。

names = ['loss', 'avg', 'rate', 'seconds', 'images']
result = pd.read_csv(data_path, skiprows=[x for x in range(lines) if (x<lines*1.0/((end_ite - start_ite)*1.0)*igore or x%step!=9)], error_bad_lines=\
False, names=names)
result.head()
for name in names:
    result[name] = result[name].str.split(' ').str.get(1)

result.head()
print(result.tail())

for name in names:
    result[name] = pd.to_numeric(result[name])
result.dtypes
print('result',result['avg'].values)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)

x_num = len(result['avg'].values)
tmp = (end_ite-start_ite - igore)/(x_num*1.0)
x = []
for i in range(x_num):
    x.append(i*tmp + start_ite + igore)
#print(x)
print('total = %d\n' %x_num)
print('start = %d, end = %d\n' %(x[0], x[-1]))

ax.plot(x, result['avg'].values, label='avg_loss')
#ax.plot(result['loss'].values, label='loss')
plt.yticks(y_ticks)#如果不想自己设置纵坐标，可以注释掉。
plt.grid()
ax.legend(loc = 'best')
ax.set_title('The loss curves')
ax.set_xlabel('batches')
fig.savefig(result_path)
#fig.savefig('loss')

• IOU loss

#!/usr/bin/python
#coding=utf-8

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

#根据log_iou修改行数
lines = 952812
step = 5000
start_ite = 0
end_ite = 50200
igore = 100
data_path =  'train_log_iou.txt' #train_log_iou的路径。
result_path = 'avg_iou' #保存结果的路径。

names = ['Region Avg IOU', 'Class', 'Obj', 'No Obj', '.5_Recall', '.7_Recall', 'count']
#result = pd.read_csv('log_iou.txt', skiprows=[x for x in range(lines) if (x%10==0 or x%10==9)]\
result = pd.read_csv(data_path, skiprows=[x for x in range(lines) if (x<lines*1.0/((end_ite - start_ite)*1.0)*igore or x%step!=0)]\
, error_bad_lines=False, names=names)
result.head()

for name in names:
    result[name] = result[name].str.split(': ').str.get(1)
result.head()
result.tail()
for name in names:
    result[name] = pd.to_numeric(result[name])
result.dtypes

x_num = len(result['Region Avg IOU'].values)
tmp = (end_ite-start_ite - igore)/(x_num*1.0)
x = []
for i in range(x_num):
    x.append(i*tmp + start_ite + igore)
#print(x)
print('total = %d\n' %x_num)
print('start = %d, end = %d\n' %(x[0], x[-1]))

fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot(x, result['Region Avg IOU'].values, label='Region Avg IOU')
#ax.plot(result['Avg Recall'].values, label='Avg Recall')
plt.grid()
ax.legend(loc='best')
ax.set_title('The Region Avg IOU curves')
ax.set_xlabel('batches')
fig.savefig(result_path)