YOLOV7训练自己的数据集

 本文以训练NWPU VHR-10数据集为例，NWPU VHR-10遥感数据集是由西北工业大学公布的用于遥感图像目标检测的公开数据集，包含10类地物目标共800张遥感图像，具体有airplane、ship 、storage tank 、baseball diamond、tennis court、basketball court、ground track field、harbor、bridge、vehicle等十种类别。

一，在官方的开源的下载YOLOV7，一般下载zip压缩包即可，下载链接为GitHub - WongKinYiu/yolov7: Implementation of paper - YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors

二，接下来则是配置环境，按照下载的文件中的requirements.txt里面提到的就行，这一块可以去线上搜搜配置的流程，我就不过多赘述了，如果实在不想自己动手，就去咸鱼上十来块钱找人配置一下，保证又快又好。有个小建议，很多时候下载一些库函数很慢，因此大家可以使用一些镜像源。

三，数据集的制作（这一块为了偷懒就直接用自己上一篇博客中YOLOv5的数据及制作流程），对于大多数的数据集来说一般是VOC格式的，在这里也就逐步的给大家具体介绍一下如何去做符合官方代码训练的数据集格式。一般来讲，我们下载的数据集的分布大概如下图所示：

其中Annotations里面存的是.xml的标签文件，而JPEGImages则是存放的图片，并且标签名和图片名一一对应，接下来，我们在打开的文件夹中先建立一个新的文件夹data1，如下所示：

之后直接将下载好的VOC格式的数据集放在此文件夹下（其中ImageSets可有可无，不影响）。

接下来我们需要做两个步骤，1：将VOC数据集转为YOLO格式的数据集。2：对转换好的数据集进行划分。

这里我们使用两个脚本，分别是用来转换标签的类别和对数据集进行划分。

1.标签类别转换，这里要记住classes里面写的是数据集中的类别。

import xml.etree.ElementTree as ET
import os, cv2
import numpy as np
from os import listdir
from os.path import join

classes = ['airplane','baseball diamond','basketball court','bridge','ship','storage tank','ground track field','harbor','vehicle','tennis court']#这个地方是数据集的类别名称

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(xmlpath, xmlname):
    with open(xmlpath, "r", encoding='utf-8') as in_file:
        txtname = xmlname[:-4] + '.txt'
        txtfile = os.path.join(txtpath, txtname)
        tree = ET.parse(in_file)
        root = tree.getroot()
        filename = root.find('filename')
        img = cv2.imdecode(np.fromfile('{}/{}.{}'.format(imgpath, xmlname[:-4], postfix), np.uint8), cv2.IMREAD_COLOR)
        h, w = img.shape[:2]
        res = []
        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes:
                classes.append(cls)
            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))
            bb = convert((w, h), b)
            res.append(str(cls_id) + " " + " ".join([str(a) for a in bb]))
        if len(res) != 0:
            with open(txtfile, 'w+') as f:
                f.write('\n'.join(res))


if __name__ == "__main__":
    postfix = 'jpg'
    imgpath = 'VOCdevkit/JPEGImages'
    xmlpath = 'VOCdevkit/Annotations'
    txtpath = 'VOCdevkit/txt'
    
    if not os.path.exists(txtpath):
        os.makedirs(txtpath, exist_ok=True)
    
    list = os.listdir(xmlpath)
    error_file_list = []
    for i in range(0, len(list)):
        try:
            path = os.path.join(xmlpath, list[i])
            if ('.xml' in path) or ('.XML' in path):
                convert_annotation(path, list[i])
                print(f'file {list[i]} convert success.')
            else:
                print(f'file {list[i]} is not xml format.')
        except Exception as e:
            print(f'file {list[i]} convert error.')
            print(f'error message:\n{e}')
            error_file_list.append(list[i])
    print(f'this file convert failure\n{error_file_list}')
    print(f'Dataset Classes:{classes}')

2，数据集比例划分

import os, shutil
from sklearn.model_selection import train_test_split

val_size = 0.1
test_size = 0.2
postfix = 'jpg'
imgpath = 'VOCdevkit/JPEGImages'
txtpath = 'VOCdevkit/txt'

os.makedirs('images/train', exist_ok=True)
os.makedirs('images/val', exist_ok=True)
os.makedirs('images/test', exist_ok=True)
os.makedirs('labels/train', exist_ok=True)
os.makedirs('labels/val', exist_ok=True)
os.makedirs('labels/test', exist_ok=True)

listdir = os.listdir(txtpath)
train, test = train_test_split(listdir, test_size=test_size, shuffle=True, random_state=0)
train, val = train_test_split(train, test_size=val_size, shuffle=True, random_state=0)

for i in train:
    shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'images/train/{}.{}'.format(i[:-4], postfix))
    shutil.copy('{}/{}'.format(txtpath, i), 'labels/train/{}'.format(i))

for i in val:
    shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'images/val/{}.{}'.format(i[:-4], postfix))
    shutil.copy('{}/{}'.format(txtpath, i), 'labels/val/{}'.format(i))

for i in test:
    shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'images/test/{}.{}'.format(i[:-4], postfix))
    shutil.copy('{}/{}'.format(txtpath, i), 'labels/test/{}'.format(i))

之后把这两个脚本放在data1文件夹下，并依次运行，记住不要放错目录级别。

依次运行后，data1文件夹下的格式便有了变化，其中新生成的images是划分后的图片，labels则是划分后并转换好的标签文件。

四，对文件进行配置工作。

1，首先，把deploy和training文件夹下面的yolov7.yaml文件中的nc都改为你训练的数据集的类别数目，我们这里数据集一共有十个目标类别，因此改为10.

2，我们需要在data目录下新建一个.yaml文件，这个文件是对我们自己数据集的信息进行描述，其中包括数据集的类别数以及不同类别的名称，还有数据集图片存放的位置。因为我使用的是NWPU数据集，因此就新建一个文件为NWPU.yaml，下图所示是他的全部内容。

3，接下来我们开始对训练脚本进行配置，找到根目录下面的train.py文件，首先我们需要下载一下不同版本的模型对应的权重，我这里就是使用正常的yolov7，这里下载路径直接去README.md里面找就行，下载好了放在根目录下即可，然后把权重路径改为你的权重存放路径就可以了。其次，还需要配置--cfg，这里存放的是模型的参数文件，这里尽量要写的完整，接下来--data则为我们自己配置好的数据集的信息yaml文件路径，然后一些其他的就不赘述了，如果显存不够，可以把batchsize改小，最小为2，尽量不要为1，如果还是跑不动，可以把图像尺寸也改小，但是由于YOLOv7的最大下采样倍数是32，因此图像尺寸要是32的整数倍。

如果还有什么问题，可以在评论区发出来，或者私聊我。