YOLOv7訓(xùn)練自己的數(shù)據(jù)集整個(gè)過程主要包括：環(huán)境安裝—制作數(shù)據(jù)集—模型訓(xùn)練—模型測試—模型推理

一、準(zhǔn)備深度學(xué)習(xí)環(huán)境

本人的筆記本電腦系統(tǒng)是：Windows10

首先下載YOLOv7的代碼，手動(dòng)下載zip或是git clone 遠(yuǎn)程倉庫，本人下載的是YOLOv7的0.1版本代碼：

https://github.com/WongKinYiu/yolov7

代碼文件夾中會(huì)有requirements.txt文件，里面描述了所需要的安裝包。

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

二、 準(zhǔn)備自己的數(shù)據(jù)集

本人標(biāo)注的數(shù)據(jù)格式是VOC，而YOLOv7能夠直接使用的是YOLO格式的數(shù)據(jù)，因此下面將介紹如何將自己的數(shù)據(jù)集轉(zhuǎn)換成可以直接讓YOLOv7進(jìn)行使用。

1、創(chuàng)建數(shù)據(jù)集

在YOLOv7文件夾中的data目錄下創(chuàng)建mydata文件夾（名字可以自定義），目錄結(jié)構(gòu)如下，將之前l(fā)abelImg標(biāo)注好的xml文件和圖片放到對應(yīng)目錄下

mydata

…images # 存放圖片

…xml # 存放圖片對應(yīng)的xml文件

…dataSet #之后會(huì)在Main文件夾內(nèi)自動(dòng)生成train.txt，val.txt，test.txt和trainval.txt四個(gè)文件，存放訓(xùn)練集、驗(yàn)證集、測試集圖片的名字（無后綴.jpg）

示例如下：

mydata文件夾下內(nèi)容如下：

image為VOC數(shù)據(jù)集格式中的JPEGImages，內(nèi)容如下：

xml文件夾下面為.xml文件（標(biāo)注工具采用labelImage），內(nèi)容如下：

dataSet 文件夾下面存放訓(xùn)練集、驗(yàn)證集、測試集的劃分，通過腳本生成，可以創(chuàng)建一個(gè)split_train_val.py文件，代碼內(nèi)容如下：

# coding:utf-8
 
import os
import random
import argparse
 
parser = argparse.ArgumentParser()
# xml文件的地址，根據(jù)自己的數(shù)據(jù)進(jìn)行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='xml', type=str, help='input xml label path')
# 數(shù)據(jù)集的劃分，地址選擇自己數(shù)據(jù)下的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] + '
'
  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()

運(yùn)行代碼后，在dataSet 文件夾下生成下面四個(gè)txt文檔：

三個(gè)txt文件里面的內(nèi)容如下：

2、轉(zhuǎn)換數(shù)據(jù)格式

接下來準(zhǔn)備labels，把數(shù)據(jù)集格式轉(zhuǎn)換成yolo_txt格式，即將每個(gè)xml標(biāo)注提取bbox信息為txt格式，每個(gè)圖像對應(yīng)一個(gè)txt文件，文件每一行為一個(gè)目標(biāo)的信息，包括class, x_center, y_center, width, height格式。格式如下：

創(chuàng)建voc_label.py文件，將訓(xùn)練集、驗(yàn)證集、測試集生成label標(biāo)簽（訓(xùn)練中要用到），同時(shí)將數(shù)據(jù)集路徑導(dǎo)入txt文件中，代碼內(nèi)容如下：

# -*- 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
    # 標(biāo)注越界修正
    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]) + '
')
 
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('mydata/%s.txt' % (image_set), 'w')
  for image_id in image_ids:
    list_file.write(abs_path + '/mydata/images/%s.jpg
' % (image_id))
    convert_annotation(image_id)
  list_file.close()

3、配置文件

1）數(shù)據(jù)集的配置

在YOLOv7目錄下的data文件夾下新建一個(gè)mydata.yaml文件（可以自定義命名），用來存放訓(xùn)練集和驗(yàn)證集的劃分文件（train.txt和val.txt）

這兩個(gè)文件是通過運(yùn)行voc_label.py代碼生成的，然后是目標(biāo)的類別數(shù)目和具體類別列表，mydata.yaml內(nèi)容如下：

2) 選擇一個(gè)你需要的模型

在YOLOv7目錄下的cfg/deploy文件夾下是模型的配置文件，這邊提供yolov7、yolov7-d6、yolov7-e6、yolov7-e6e、yolov7x等多個(gè)版本，假設(shè)采用yolov7x.yaml，只用修改一個(gè)參數(shù)，把nc改成自己的類別數(shù),需要取整（可選） 如下：

至此，自定義數(shù)據(jù)集已創(chuàng)建完畢，接下來就是訓(xùn)練模型了。

三、模型訓(xùn)練

1、下載預(yù)訓(xùn)練模型

在YOLOv7的GitHub開源網(wǎng)址上下載對應(yīng)版本的模型

2、訓(xùn)練

在正式開始訓(xùn)練之前，需要對train.py進(jìn)行以下修改：

以上參數(shù)解釋如下：

epochs：指的就是訓(xùn)練過程中整個(gè)數(shù)據(jù)集將被迭代多少次,顯卡不行你就調(diào)小點(diǎn)。

batch-size：一次看完多少張圖片才進(jìn)行權(quán)重更新，梯度下降的mini-batch,顯卡不行你就調(diào)小點(diǎn)。

cfg：存儲(chǔ)模型結(jié)構(gòu)的配置文件

data：存儲(chǔ)訓(xùn)練、測試數(shù)據(jù)的文件

img-size：輸入圖片寬高,顯卡不行你就調(diào)小點(diǎn)。

之后運(yùn)行訓(xùn)練命令如下：

python train.py --img 640 --batch 32 --epoch 300 --data data/mydata.yaml --cfg cfg/deploy/yolov7x.yaml --weights weights/yolov7x.pt --device '0' 

四、模型測試

評估模型好壞就是在有標(biāo)注的測試集或者驗(yàn)證集上進(jìn)行模型效果的評估，在目標(biāo)檢測中最常使用的評估指標(biāo)為mAP。在test.py文件中指定數(shù)據(jù)集配置文件和訓(xùn)練結(jié)果模型，如下：

通過下面的命令進(jìn)行模型測試：

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

模型測試效果如下：

五、模型推理

最后，模型在沒有標(biāo)注的數(shù)據(jù)集上進(jìn)行推理，在detect.py文件中指定測試圖片和測試模型的路徑

其他參數(shù)（img_size、置信度object confidence threshold、IOU threshold for NMS）可自行修改，如下：

使用下面的命令，其中，weights使用最滿意的訓(xùn)練模型即可，source則提供一個(gè)包含所有測試圖片的文件夾路徑即可。

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

測試完畢后，每個(gè)測試圖片會(huì)在指定的inference/output輸出文件夾中生成結(jié)果圖片文件，如下：