YOLOv5训练自己的数据集

小白教学

环境配置

Anaconda创建、激活、退出、删除虚拟环境

退出环境
conda deactivate

创建环境
conda create -n zyf_test python=3.7

激活环境
source activate zyf_test

cd yolov5-master
pip install -r reqirements.txt

准备数据集（VOC格式）

启动LabelImg

(base) C:\Users\xxx>conda env list
#conda environments:

base * D:\anaconda
LabelIMG D:\anaconda\envs\LabelIMG
pytorch D:\anaconda\envs\pytorch

(base) C:\Users\fanfan>conda activate LabelIMG
(LabelIMG) C:\Users\fanfan>labelImg

创建voc_data文件夹

paper_data
…images # 存放图片
…Annotations # 存放图片对应的xml文件
…ImageSets/Main

Annotations文件夹下面为xml文件（标注工具采用labelImage），内容如下：

images为VOC数据集格式中的JPEGImages，内容如下：

创建一个split.py文件：

import os
import random

trainval_percent = 0.9  # 训练和验证集所占比例，剩下的0.1就是测试集的比例
train_percent = 0.8  # 训练集所占比例，可自己进行调整
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets\Main'
total_xml = os.listdir(xmlfilepath)

num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/Main/val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

运行代码后，在Main文件夹下生成下面四个txt文档，要放在ImageSets/Main里：

准备labels

接下来准备labels，把数据集格式转换成yolo_txt格式，即将每个xml标注提取bbox信息为txt格式（这种数据集格式成为yolo_txt格式），每个图像对应一个txt文件，文件每一行为一个目标的信息，包括类别 xmin xmax ymin ymax。

创建my_voc_label.py文件，将训练集、验证集、测试集生成label标签（训练中要用到），同时将数据集路径导入txt文件中：

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

#sets=[('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
#sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
sets=[('train'), ('val'), ('test')]

#classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
classes = ["t1"]  # 改成自己的类别


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('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').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))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

wd = getcwd()  #getcwd()会将当前工作目录的绝对路径获取到参数buffer

for image_set in sets:
    if not os.path.exists('labels'):
        os.makedirs('labels')
    image_ids = open('ImageSets/Main/%s.txt'%(image_set)).read().strip().split()
    list_file = open('%s.txt'%(image_set), 'w')
    for image_id in image_ids:
        list_file.write('voc_data/images/%s.jpg\n'%(image_id))
        convert_annotation(image_id)
    list_file.close()

# 路徑拼接 暂时用不到
#os.system("cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt > train.txt")
#os.system("cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt > train.all.txt")

运行后会生成如下labels文件夹和三个包含数据集的txt文件，其中labels中为不同图像的标注文件，train.txt等txt文件为划分后图像所在位置的绝对路径，如train.txt就含有所有训练集图像的绝对路径。

配置文件

数据集的配置

在yolov5目录下的data文件夹下新建一个t1.yaml文件（可以自定义命名），用来存放训练集和验证集的划分文件（train.txt和val.txt），这两个文件是通过运行my_voc_label.py代码生成的，然后是目标的类别数目和具体类别列表，t1.yaml内容如下：

模型训练

下载预训练模型

源码中在yolov5目录下的weights文件夹下提供了下载四种预训练模型的脚本----download_weights.sh，执行这个shell脚本就可以下载。（也可百度网盘下载，但是不一定是v3.0版本最新的预训练模型，如果不是v3.0版本最新的预训练模型，在训练时会报错）

训练

train.py 修改data，epoches，batchsize
imgsize电脑不行会报错，可以改小一点

参数解释：
epochs：指的就是训练过程中整个数据集将被迭代多少次,显卡不行你就调小点。
batch-size：一次看完多少张图片才进行权重更新，梯度下降的mini-batch,显卡不行你就调小点。
cfg：存储模型结构的配置文件
data：存储训练、测试数据的文件
img-size：输入图片宽高,显卡不行你就调小点。
rect：进行矩形训练
resume：恢复最近保存的模型开始训练
nosave：仅保存最终checkpoint
notest：仅测试最后的epoch
evolve：进化超参数
bucket：gsutil bucket
cache-images：缓存图像以加快训练速度
weights：权重文件路径
name： 重命名results.txt to results_name.txt
device：cuda device, i.e. 0 or 0,1,2,3 or cpu
adam：使用adam优化
multi-scale：多尺度训练，img-size +/- 50%
single-cls：单类别的训练集

训练好的模型会被保存在yolov5目录下的runs/train/exp13/weights/last.pt和best.pt。详细训练数据保存在runs/train/exp13/results.txt文件中

训练过程可视化

runs/train/exp13文件夹内。

模型测试

评估模型好坏就是在有标注的测试集或者验证集上进行模型效果的评估，在目标检测中最常使用的评估指标为mAP。在test.py文件中指定数据集配置文件和训练结果模型，如下：
修改weights和data

模型推理

最后，模型在没有标注的数据集上进行推理，在detect.py文件中指定测试图片和测试模型的路径，其他参数（img_size、置信度object confidence threshold、IOU threshold for NMS）可自行修改，如下：

其中，测试图片在data/images文件夹中
结果图片保存在runs/detect/exp…文件夹中