最近在做一個用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']
- 將相關的路徑修改爲自己的路徑
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"
跟圖片測試一眼,雖然能夠成功運行,但是檢測的效果很差,另外增加了一些程序運行的信息,發現檢測速度也不是很好。後來修改代碼,使用幀交替雙線程的方法進行測試,發現不僅速度快了,檢測效果也變好了,具體是什麼原因導致的,目前我也不清楚。幀交替雙線程的視頻檢測代碼將在下一篇博客中進行敘述。