YOLOv3+Tensorflow訓練自己的數據

最近在做一個用YOLOv3進行目標檢測的項目,根據網上各位大神的博客,結合自己做項目的過程做了一個筆記,方便自己以後回顧實驗過程,也給大家做個參考。
實驗用的程序主代碼來自於github上一位大神程序猿,這裏給出代碼的地址
YOLOv3-tensorflow大神源碼

一、製作數據集

根據代碼來看,訓練和測試使用的數據格式需要按照PascalVOC的數據格式來,所以首先我們需要進行數據集的製作。

1 獲取數據

根據實驗需要,我將錄製的視頻文件分爲訓練視頻和測試視頻,然後按幀進行截取,獲得訓練和測試用的圖片,因爲實驗的保密性,在此不能明具體的內容。將視頻按幀截取爲圖片並保存使用的是下面這段代碼:

import cv2
import glob
import os
from datetime import datetime

def video_to_frames(path):
    """
    輸入:path(視頻文件的路徑)
    """
    # VideoCapture視頻讀取類
    videoCapture = cv2.VideoCapture()
    videoCapture.open(path)
    for i in range(int(frames)):
        ret, frame = videoCapture.read()
        cv2.imwrite("E:\dataset\images\train%d.jpg" % (i), frame)
    return

if __name__ == '__main__':
    video_to_frames("E:\dataset\video\train.mp4")
    print("SUCCEED !!!")

2 標記圖片

根據PascalVOC數據集的需要,使用Labelimg工具對圖片進行標註,標註後會生成XML文件,如下圖所示:
在這裏插入圖片描述

2 按照PascalVOC數據集的格式整理自己的數據

這次實驗我分別建立了VOCTrainval和VOCTest兩個數據文件,分別用於訓練和測試,大家也可以不分開,後面進行訓練和測試數據劃分就行了,兩個文件夾都按照PascalVOC的格式建立。PascalVOC數據集包含了5個部分,在實驗中我們只需要用到一下三個文件夾:
1) Annatations文件夾
文件夾存放的是xml格式的標籤文件,每個xml文件都對應於JPEGImages文件夾的一張圖片。
2)JPEGImages文件夾
文件夾裏包含了訓練圖片或測試圖片。
3)ImageSets文件夾
該文件夾裏原有三個子文件夾,但實驗中我們僅需要使用Main文件夾裏面的信息,存放的是圖像物體識別的數據,有train.txt, val.txt ,trainval.txt.這三個文件(VOCTrainval文件夾下)或者test.txt 文件(VOCTest)。這幾個文件我們後面會生成。
按照要求,將自己的圖片放入JPEGImages文件夾,將標註信息xml文件放入Annatations文件夾:
在這裏插入圖片描述
在這裏插入圖片描述

3 劃分訓練集和測試集

訓練時要有測試集和訓練集,如果在製作數據集的時候沒有像我一樣進行區分,那麼在這裏就需要使用代碼將數據進行劃分,放在ImageSets\Main文件夾下。代碼如下,至於訓練驗證集和測試集的劃分比例,以及訓練集和驗證集的劃分比例,根據自己的數據情況決定。將下面的代碼放入split.py中:

import os  
import random   
  
xmlfilepath=r'E:/tensorflow-yolov3-master/VOCData/VOCTrainVal/Annotations/'  #xml文件的路徑
saveBasePath=r'E:/tensorflow-yolov3-master/VOCData/VOCTrainVal/ImageSets/' #生成的txt文件的保存路徑
  
trainval_percent=0.9 
train_percent=0.8  
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)    
  
print("train and val size",tv)  
print("traub suze",tr)  
ftrainval = open(os.path.join(saveBasePath,'Main/trainval.txt'), 'w')    
ftest = open(os.path.join(saveBasePath,'Main/test.txt'), 'w')    
ftrain = open(os.path.join(saveBasePath,'Main/train.txt'), 'w')    
fval = open(os.path.join(saveBasePath,'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:    
            ftrain.write(name)    
        else:    
            fval.write(name)    
    else:    
        ftest.write(name)    
    
ftrainval.close()    
ftrain.close()    
fval.close()    
ftest .close() 

直接運行split.py,或者對路徑和劃分比例稍做修改後分別對VOCTrainval和VOCTest文件夾中的數據運行即可,得到ImageSets\Main\文件下的幾個txt文件。到此數據集製作的第一個階段完成!

二、在data/dataset文件夾下生成數據的描述文件

根據作者在github上的聲明,該程序在訓練時還需要兩個對圖片信息的描述文件voc_train.txt和voc_test.txt,格式如下
在這裏插入圖片描述
作者給出了生成這兩個文件的代碼scripts\voc_annotation.py。需要做部分的修改:
在這裏插入圖片描述
1)將classes修改爲自己的類別

classes = ['類別1','類別2','類別3','類別4']
  1. 將相關的路徑修改爲自己的路徑
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 將相關路徑中default的路徑修改爲自己的路徑
    parser.add_argument("--data_path", default="E:/tensorflow-yolov3-master/VOCData/")
    # 將default的路徑更改爲自己的voc_train.txt存放的位置
    parser.add_argument("--train_annotation", default="../data/dataset/voc_train.txt")
     # 將default的路徑更改爲自己的voc_test.txt存放的位置
    parser.add_argument("--test_annotation",  default="../data/dataset/voc_test.txt")
    flags = parser.parse_args()

    if os.path.exists(flags.train_annotation):os.remove(flags.train_annotation)
    if os.path.exists(flags.test_annotation):os.remove(flags.test_annotation)

    # 更改訓練集和測試集的相對路徑,由於本實驗的訓練數據只有一個路徑,所以註釋掉了num2
    num1 = convert_voc_annotation(os.path.join(flags.data_path, 'VOCTrainVal'), 'trainval', flags.train_annotation, False)
    #num2 = convert_voc_annotation(os.path.join(flags.data_path, 'VOCTrainVal'), 'trainval', flags.train_annotation, False)
    num3 = convert_voc_annotation(os.path.join(flags.data_path, 'VOCTest'),  'test', flags.test_annotation, False)
    print('=> The number of image for train is: %d\tThe number of image for test is:%d' %(num1, num3))

做了以上更改後直接運行該文件即可。這時候data\dataset文件加下將多出兩個文件:
在這裏插入圖片描述
查看其中一個文件的內容如下:
在這裏插入圖片描述

三、修改其他文件

1 修改voc_names

打開data/classes/文件夾中的voc_names文件,將其中的類別名稱改爲自己的類別。
在這裏插入圖片描述

2 修改配置文件

打開core\文件夾下的config.py文件
在這裏插入圖片描述
根據自己的文件路徑,進行適當修改,當然,若按照我上述步驟來的,路徑應該是沒有什麼需要修改的地方的,那些參數也自己按照訓練的過程慢慢調就行了。

四、訓練

接下來就是訓練啦,直接運行train.py就行了,一般來說是不會有什麼問題的。

五、測試

訓練完成之後,終於到了最激動人心的時刻了,進行模型的測試!

1 修改config文件

修改__C.TEST.ANNOT_PATH爲自己voc_test.txt文件的路徑
修改__C.TEST.WEIGHT_FILE爲自己需要測試的模型的路徑
__C.TEST.INPUT_SIZE根據需要自行選擇,也可以不做修改

# TEST options
__C.TEST                        = edict()

__C.TEST.ANNOT_PATH             = "./data/dataset/voc_test.txt"  #修改爲自己voc_test.txt文件的路徑
__C.TEST.BATCH_SIZE             = 6
__C.TEST.INPUT_SIZE             = 544  #自行決定是否修改,不修改也沒有影響
__C.TEST.DATA_AUG               = False
__C.TEST.WRITE_IMAGE            = True
__C.TEST.WRITE_IMAGE_PATH       = "./data/detection/"
__C.TEST.WRITE_IMAGE_SHOW_LABEL = True
__C.TEST.WEIGHT_FILE            = "./checkpoint/yolov3.ckpt" # 修改爲自己模型的路徑
__C.TEST.SHOW_LABEL             = True
__C.TEST.SCORE_THRESHOLD        = 0.3
__C.TEST.IOU_THRESHOLD          = 0.45

2 執行evaluate.py文件

這裏沒什麼需要修改的,我後來是有增加一個保存檢測後圖片的代碼,因爲作者給的圖片測試代碼效果太差了,這個後面再講。

3 修改mAP\extra\中的class_list.txt文件

在這裏插入圖片描述
打開class_list.txt文件,將內容修改爲自己的類別。

4 計算mAP和各類AP

1)若類別名稱中有空格,需要先執行mAP\extra\文件夾中的remove_space.py或者rename_class.py,將空格轉換成短橫‘-’。若類別名中沒有空格則忽略這一步。
在這裏插入圖片描述
2)執行mAP\main.py,得出各類AP以及mAP等,其結果保存在mAP\results\目錄下。
在這裏插入圖片描述
在這裏插入圖片描述

六、圖片測試和視頻測試

作者給出的測試代碼需要將模型另外保存爲.pb的形式,所以需要先對模型進行轉換。

1 生成YOLOv3.pb文件

修改freeze_graph.py
1)將pb_file修改爲自己希望.pb模型保存的路徑

pb_file = "./model_pb/yolov3.pb"

2)將ckpt_file修改爲自己的.ckpt模型的路徑

ckpt_file = "./checkpoint/yolov3_test_loss=7.8478.ckpt-62"

然後直接運行freeze_graph.py文件。

2 圖片測試

修改image_demo.py
將pb_file修改爲自己模型的路徑
將video_path修改爲自己圖片的路徑

pb_file         = "./model_pb/yolov3.pb"
video_path      = "./docs/images/test1.jpg"

運行image_demo.py
不知道爲啥,效果很差,但是在上一步的測試中mAP已經達到了87%,效果不應該會有這麼差的,所以我就直接在evaluate.py中增加了幾行代碼,將evaluate.py中生成的檢測後的圖片全部保存了下來,果然,檢測的效果是很好的,增加的幾行代碼如下:

if self.write_image:
                    image = utils.draw_bbox(image, bboxes_pr, show_label=self.show_label)
                    cv2.imwrite(self.write_image_path+image_name, image)
                    # 下面三行爲新增代碼
                    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                    image = Image.fromarray(image)
                    image.save(output_path+image_name)

3 視頻測試

跟上一步一樣修改video_demo.py
將pb_file修改爲自己模型的路徑
將video_path修改爲自己圖片的路徑

pb_file         = "./model_pb/yolov3.pb"
video_path      = "./docs/images/test.mp4"

跟圖片測試一眼,雖然能夠成功運行,但是檢測的效果很差,另外增加了一些程序運行的信息,發現檢測速度也不是很好。後來修改代碼,使用幀交替雙線程的方法進行測試,發現不僅速度快了,檢測效果也變好了,具體是什麼原因導致的,目前我也不清楚。幀交替雙線程的視頻檢測代碼將在下一篇博客中進行敘述。

本人目標檢測初學者,博客中或許有一些不盡完善的地方,歡迎大家提出來一起交流討論。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章