目錄

一、準備深度學習環境

二、 準備自己的資料集

1、創建資料集

2、轉換資料格式

3、組態檔

三、模型訓練

1、下載預訓練模型

2、訓練

四、模型測驗

五、模型推理

YOLOv5訓練自己的資料集整個程序主要包括：環境安裝----制作資料集----模型訓練----模型測驗----模型推理

一、準備深度學習環境

本人的筆記本電腦系統是：Windows10
首先進入YOLOv5開源網址 ，手動下載zip或是git clone 遠程倉庫，本人下載的是YOLOv5的5.0版本代碼，代碼檔案夾中會有requirements.txt檔案，里面描述了所需要的安裝包，

本文最終安裝的pytorch版本是1.8.1，torchvision版本是0.9.1，python是3.7.10，其他的依賴庫按照requirements.txt檔案安裝即可，

二、 準備自己的資料集

本人在訓練YOLOv5時，選擇的資料格式是VOC，因此下面將介紹如何將自己的資料集轉換成可以直接讓YOLOv5進行使用，

1、創建資料集

在YOLOv5檔案夾中的data目錄下創建mydata檔案夾（名字可以自定義），目錄結構如下，將之前labelImg標注好的xml檔案和圖片放到對應目錄下
mydata
…images # 存放圖片
…xml # 存放圖片對應的xml檔案
…dataSet #之后會在Main檔案夾內自動生成train.txt，val.txt，test.txt和trainval.txt四個檔案，存放訓練集、驗證集、測驗集圖片的名字（無后綴.jpg）
示例如下：
mydata檔案夾下內容如下：

• image為VOC資料集格式中的JPEGImages，內容如下：

• xml檔案夾下面為.xml檔案（標注工具采用labelImage），內容如下：

• dataSet 檔案夾下面存放訓練集、驗證集、測驗集的劃分，通過腳本生成，可以創建一個split_train_val.py檔案，代碼內容如下：

# coding:utf-8

import os
import random
import argparse

parser = argparse.ArgumentParser()
# xml檔案的地址，根據自己的資料進行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='xml', type=str, help='input xml label path')
# 資料集的劃分，地址選擇自己資料下的ImageSets/Main
parser.add_argument('--txt_path', default='dataSet', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0
train_percent = 0.9
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()

• 運行代碼后，在dataSet 檔案夾下生成下面四個txt檔案：

• 三個txt檔案里面的內容如下：

2、轉換資料格式

接下來準備labels，把資料集格式轉換成yolo_txt格式，即將每個xml標注提取bbox資訊為txt格式，每個影像對應一個txt檔案，檔案每一行為一個目標的資訊，包括class, x_center, y_center, width, height格式，格式如下：

• 創建voc_label.py檔案，將訓練集、驗證集、測驗集生成label標簽（訓練中要用到），同時將資料集路徑匯入txt檔案中，代碼內容如下：

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd

sets = ['train', 'val', 'test']
classes = ["a", "b"]   # 改成自己的類別
abs_path = os.getcwd()
print(abs_path)

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h

def convert_annotation(image_id):
    in_file = open('data/mydata/xml/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('data/mydata/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').text
        difficult = obj.find('Difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_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 = w
        if b4 > h:
            b4 = h
        b = (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('data/mydata/labels/'):
        os.makedirs('data/mydata/labels/')
    image_ids = open('data/mydata/dataSet/%s.txt' % (image_set)).read().strip().split()
    list_file = open('paper_data/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write(abs_path + '/mydata/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

3、組態檔

1）資料集的配置
在yolov5目錄下的data檔案夾下新建一個mydata.yaml檔案（可以自定義命名），用來存放訓練集和驗證集的劃分檔案（train.txt和val.txt），這兩個檔案是通過運行voc_label.py代碼生成的，然后是目標的類別數目和具體類別串列，mydata.yaml內容如下：

2) 選擇一個你需要的模型
在yolov5目錄下的model檔案夾下是模型的組態檔，這邊提供s、m、l、x版本，逐漸增大（隨著架構的增大，訓練時間也是逐漸增大），假設采用yolov5x.yaml，只用修改一個引數，把nc改成自己的類別數,需要取整（可選） 如下：

至此，自定義資料集已創建完畢，接下來就是訓練模型了，

三、模型訓練

1、下載預訓練模型

在YOLOv5的GitHub開源網址上下載對應版本的模型

2、訓練

在正式開始訓練之前，需要對train.py進行以下修改：

以上引數解釋如下：
epochs：指的就是訓練程序中整個資料集將被迭代多少次,顯卡不行你就調小點，
batch-size：一次看完多少張圖片才進行權重更新，梯度下降的mini-batch,顯卡不行你就調小點，
cfg：存盤模型結構的組態檔
data：存盤訓練、測驗資料的檔案
img-size：輸入圖片寬高,顯卡不行你就調小點，

之后運行訓練命令如下：

python train.py --img 640 --batch 32 --epoch 300 --data data/mydata.yaml --cfg models/yolov5x.yaml --weights weights/yolov5x.pt --device '0,1' 

四、模型測驗

評估模型好壞就是在有標注的測驗集或者驗證集上進行模型效果的評估，在目標檢測中最常使用的評估指標為mAP，在test.py檔案中指定資料集組態檔和訓練結果模型，如下：

通過下面的命令進行模型測驗：

python test.py  --data data/mydata.yaml --weights runs/exp1/weights/best.pt --augment

模型測驗效果如下：

五、模型推理

最后，模型在沒有標注的資料集上進行推理，在detect.py檔案中指定測驗圖片和測驗模型的路徑，其他引數（img_size、置信度object confidence threshold、IOU threshold for NMS）可自行修改，如下：

使用下面的命令，其中，weights使用最滿意的訓練模型即可，source則提供一個包含所有測驗圖片的檔案夾路徑即可，

 python detect.py --weights runs/exp1/weights/best.pt --source inference/images/ --device 0,1

測驗完畢后，每個測驗圖片會在指定的inference/output輸出檔案夾中生成結果圖片檔案，如下：

本人訓練的資料集是口罩資料集，檢測后的效果如下圖所示：