YOLOv7訓練自己的數據集整個過程主要包括：環境安裝—制作數據集—模型訓練—模型測試—模型推理

一、準備深度學習環境

本人的筆記本電腦系統是：Windows10

首先下載YOLOv7的代碼，手動下載zip或是git clone 遠程倉庫，本人下載的是YOLOv7的0.1版本代碼：

https://github.com/WongKinYiu/yolov7

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

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

二、 準備自己的數據集

本人標注的數據格式是VOC，而YOLOv7能夠直接使用的是YOLO格式的數據，因此下面將介紹如何將自己的數據集轉換成可以直接讓YOLOv7進行使用。

1、創建數據集

在YOLOv7文件夾中的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] + '
'
  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]) + '
')
 
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）數據集的配置

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

這兩個文件是通過運行voc_label.py代碼生成的，然后是目標的類別數目和具體類別列表，mydata.yaml內容如下：

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

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

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

三、模型訓練

1、下載預訓練模型

在YOLOv7的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 cfg/deploy/yolov7x.yaml --weights weights/yolov7x.pt --device '0' 

四、模型測試

評估模型好壞就是在有標注的測試集或者驗證集上進行模型效果的評估，在目標檢測中最常使用的評估指標為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輸出文件夾中生成結果圖片文件，如下：