YOLOv5自定義數據集訓練

簡介

本文介紹如何在自己的VOC格式數據集上訓練YOLO5目標檢測模型。

VOC數據集格式

首先，先來了解一下Pascal VOC數據集的格式，該數據集油5個部分組成，文件組織結構如下，目前主要的是VOC2007和VOC2012.

- VOC- JPEGImages- 1.jpg- 2.jpg- ...- Annotations- 1.xml- 2.xml- ...- ImageSets- Main- train.txt- val.txt- test.txt- trainval.txt- ...- SegmentationClass- SegmentationObject

第一個文件夾JPEGImages為所有的圖像，也就是說，訓練集、驗證集和測試集需要自己劃分；Annotations為JPEGImages文件夾中每個圖片對應的標注，xml格式文件，文件名與對應圖像相同；ImageSets主要的子文件夾為Main，其中有四個文本文件，為訓練集、驗證集、測試集和訓練驗證集的圖片文件名；SegmentationClass和SegmentationObject文件夾存放分割的結果圖，前者為語義分割，后者為實例分割。

上述xml標注文件，格式如下。對其具體標注解釋。

<annotation><folder>down</folder> # 圖片所處文件夾<filename>1.jpg</filename> # 圖片文件名及后綴<path>./savePicture/train_29635.jpg</path> # 存放路徑<source> #圖源信息<database>Unknown</database> </source><size> # 圖片尺寸和通道<width>640</width><height>480</height><depth>3</depth></size><segmented>0</segmented> #是否有分割label，0無1有# 圖像中包含的所有目標，一個目標一個object標簽<object><name>car</name> # 目標類別<pose>Unspecified</pose> # 目標的姿態<truncated>0</truncated> # 目標是否被部分遮擋（>15%）<difficult>0</difficult> # 是否為難以辨識的目標， 需要結合背景才能判斷出類別的物體<bndbox> # 目標邊界框信息<xmin>2</xmin><ymin>156</ymin><xmax>111</xmax><ymax>259</ymax></bndbox></object><object><name>multi_signs</name><editType /><pose>Unspecified</pose><truncated>0</truncated><difficult>0</difficult><bndbox><xmin>81</xmin><ymin>98</ymin><xmax>154</xmax><ymax>243</ymax></bndbox></object> </annotation>

也就是說，遇到這種文件格式的數據（主要特點為圖像全放在一個文件夾，標注格式如上等），將其作為VOC格式的數據集，將自己的數據集重構為VOC格式以便開源項目的處理。

自定義訓練

下載源碼

通過git clone git@github.com:ultralytics/yolov5.git將YOLOv5源碼下載到本地，本文后面的內容也可以參考官方的自定義數據集訓練教程，不同于我的教程，該教程全面包含了VOC格式和COCO格式數據集的處理方法。

此時創建虛擬環境，并通過pip install -r requirements.txt安裝依賴包，我這里測試過，最新的項目是兼容Pytorch 1.6的，1.6之前的Pytorch會有一些問題（這是基于本文發布的時候，隨著yolo5的更新對版本的要求會更高，后續只要按照官方教程配置環境即可）。

數據集處理

一般，符合VOC格式的數據集至少包含圖像和標注兩個文件夾，結構如下。我這里假定測試集是獨立的，該數據集實際上為訓練集，只需要劃分出訓練集和驗證集即可。這里建議將文件夾重命名如下，否則后續可能出現數據集加載失敗的情況。

- 根目錄- images- Annotations

下面，編寫腳本劃分數據集，split_train_val.py腳本內容如下（參考Github上的開源腳本），只需要執行python split_train_val.py --xml_path dataset_root/Annotations/ --txt_path dataset_root/anno_txt/就得到了劃分結果的文件列表，如訓練集對應的train.txt如下圖，里面與訓練圖片所有的文件名。

import os import random import argparseparser = argparse.ArgumentParser() parser.add_argument('--xml_path', type=str, help='input xml label path') parser.add_argument('--txt_path', type=str, help='output txt label path') opt = parser.parse_args()trainval_percent = 1.0 train_percent = 0.8 xmlfilepath = opt.xml_path txtsavepath = opt.txt_path total_xml = os.listdir(xmlfilepath) if not os.path.exists(txtsavepath):os.makedirs(txtsavepath)num = len(total_xml) list_index = range(num) tv = int(num * trainval_percent) tr = int(tv * train_percent) trainval = random.sample(list_index, tv) train = random.sample(trainval, tr)file_trainval = open(txtsavepath + '/trainval.txt', 'w') file_test = open(txtsavepath + '/test.txt', 'w') file_train = open(txtsavepath + '/train.txt', 'w') file_val = open(txtsavepath + '/val.txt', 'w')for i in list_index:name = total_xml[i][:-4] + '\n'if i in trainval:file_trainval.write(name)if i in train:file_train.write(name)else:file_val.write(name)else:file_test.write(name)file_trainval.close() file_train.close() file_val.close() file_test.close()

接下來，我們要做的就是每個xml標注提取bbox信息為txt格式（這種數據集格式成為yolo_txt格式），每個圖像對應一個txt文件，文件每一行為一個目標的信息，包括類別 xmin xmax ymin ymax。使用的腳本voc_label.py內容如下（注意，類別要替換為當前數據集的類別列表），在數據集根目錄（此時包含Annotations、anno_txt以及images三個文件夾的目錄）下執行該腳本，如python ../../utils/voc_label.py。

# -*- coding: utf-8 -*-import xml.etree.ElementTree as ET import os from os import getcwdsets = ['train', 'val', 'test'] classes = ['window_shielding', 'multi_signs', 'non_traffic_signs'] abs_path = os.getcwd()def convert(size, box):dw = 1. / (size[0])dh = 1. / (size[1])x = (box[0] + box[1]) / 2.0 - 1y = (box[2] + box[3]) / 2.0 - 1w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn x, y, w, hdef convert_annotation(image_id):in_file = open('Annotations/%s.xml' % (image_id))out_file = open('labels/%s.txt' % (image_id), 'w')tree = ET.parse(in_file)root = tree.getroot()size = root.find('size')w = int(size.find('width').text)h = int(size.find('height').text)for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),float(xmlbox.find('ymax').text))b1, b2, b3, b4 = b# 標注越界修正if b2 > w:b2 = wif b4 > h:b4 = hb = (b1, b2, b3, b4)bb = convert((w, h), b)out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')wd = getcwd() for image_set in sets:if not os.path.exists('labels/'):os.makedirs('labels/')image_ids = open('anno_txt/%s.txt' % (image_set)).read().strip().split()list_file = open('%s.txt' % (image_set), 'w')for image_id in image_ids:list_file.write(abs_path + '/images/%s.jpg\n' % (image_id))convert_annotation(image_id)list_file.close()

這時候，我們的目標檢測數據集就構建完成了，其內容如下，其中labels中為不同圖像的標注文件，train.txt等幾個根目錄下的txt文件為劃分后圖像所在位置的絕對路徑，如train.txt就含有所有訓練集圖像的絕對路徑。

配置文件

下面需要兩個配置文件用于模型的訓練，一個用于數據集的配置，一個用于模型的配置。

首先是數據集的配置，在根目錄下的data目錄下新建一個yaml文件，內容如下，首先是訓練集和驗證集的劃分文件，這個文件在上面一節最后生成得到了，然后是目標的類別數目和具體類別列表，這個列表務必和上一節最后voc_label.py中的一致。

train: dataset/train/train.txt val: dataset/train/val.txt# number of classes nc: 3# class names names: ['window_shielding', 'multi_signs', 'non_traffic_sign']

然后，編輯模型的配置文件，此時需要先在項目根目錄下的weights目錄下執行其中的download_weights.sh這個shell腳本來下載四種模型的權重。然后，選擇一個模型，編輯項目根目錄下models目錄中選擇的模型的配置文件，將第一個參數nc改為自己的數據集類別數即可，例如我使用yolov5x模型，則修改yolov5x.yaml文件。這里weights的下載可能因為網絡而難以進行，我也將其上傳到了百度網盤，地址給出，提取碼為vjlx。

模型訓練

此時，可以使用下面的命令進行模型的訓練，訓練日志默認保存在./runs/下，包括模型參數、Tensorboard記錄等。此時TensorBoard以已經默認打開，瀏覽器訪問效果如下圖（由于數據量很小，很快過擬合）。

python train.py --img 640 --batch 8 --epoch 300 --data ./data/ads.yaml --cfg ./models/yolov5x.yaml --weights weights/yolov5x.pt --device '0'

模型測試

接著，就是在有標注的測試集或者驗證集上進行模型效果的評估，在目標檢測中最常使用的指標為mAP。通過下面的命令進行模型測試，由于這是個比賽，測試集沒有標注，這里使用驗證集作為測試用數據，下述命令只需要指定數據集配置文件和訓練結果模型即可。

python test.py --data ./data/ads.yaml --weights ./runs/exp0/weights/best.pt --augment

不進行測試時數據增強和進行測試時數據增強（TTA）在驗證集上的表現分別如下。

Class Images Targets P R mAP@.5 mAP@.5:.95 all 400 970 0.376 0.441 0.35 0.235 Class Images Targets P R mAP@.5 mAP@.5:.95 all 400 970 0.272 0.532 0.366 0.24

模型推理

最后，模型在沒有標注的數據上進行推理，使用下面的命令（該命令中save-txt選項用于生成結果的txt標注文件，不指定則只會生成結果圖像）。其中，weights使用最滿意的實驗即可，source則提供一個包含所有測試圖片的文件夾即可。

python detect.py --weights runs/exp0/weights/best.pt --source ./dataset/test/ --device 0 --save-txt

這樣，對每個測試圖片會在默認的inference/output文件夾中生成一個同名的txt文件，按照我的需求修改了detect.py文件后，每個txt會生成一行一個目標的信息，信息包括類別序號 置信度 xcenter ycenter w h，后面四個為bbox位置，均未歸一化。如下圖。

我這里因為是一個比賽，再將這個txt處理為了json文件。不論是這里的處理代碼還是上面對detect.py修改的代碼，都可以在文末給出的Github倉庫找到。

補充說明

本文介紹了如何使用YOLOv5在自己的數據集上進行訓練，按部就班地進行了講解。該項目在YOLOv5地源碼基礎上修改完成。

總結

以上是生活随笔為你收集整理的YOLOv5自定义数据集训练的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。