yolov5模型训练自己的数据集

1. 环境

 Python>=3.8.0 

 PyTorch>=1.8

2. 下载yolov5

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

3. 数据集准备（VOC格式）

3.1 准备两个文件夹

Annotations存放标注信息，images存放图片

3. 2 图片标注

3.2.1 下载labellmg

https://github.com/tzutalin/labelImg

在predefined_classes.txt文件中更改标注类别

路径：labelImg-master/data/predefined_classes.txt

打开 Anaconda Prompt 

进入 labellmg文件夹

D(安装盘符):
cd 安装路径\labellmg

依次运行下面三条命令

conda install pyqt=5
conda install -c anaconda lxml
pyrcc5 -o libs/resources.py resources.qrc

如果前两行都运行成功，运行pyrcc5 -o libs/resources.py resources.qrc时，报错'pyrcc5' 不是内部或外部命令，也不是可运行的程序或批处理文件

pip install PyQt5-tools
pyrcc5 -o libs/resources.py resources.qrc

3.2.2 labellmg使用

python labelImg.py   #运行软件

勾选自动保存,打开目录-images 存放目录-Annotations

格式默认XML格式，点击左方边栏或者屏幕右键选择创建区块即可进行标注。

快捷键：  w - Create a rect box

                d - Next image

                a - Previous image

4. 数据集划分

4.1 划分训练集、验证集、测试集

创建程序split_train_val.py 并运行

# coding:utf-8

import os
import random
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', default='#xml文件的地址，根据自己的Annotations路径修改', type=str, help='input xml label path')
parser.add_argument('--txt_path', default='ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 0.8  # 训练集和验证集所占比例
train_percent = 0.9     # 训练集所占比例
test_percent = 0.2
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()

运行完后生成 ImagesSets\Main 文件夹，且在其下生成测试集、训练集、验证集

 4.2 XML格式转yolo_txt格式

创建程序 xml_to_yolo.py 并运行，注意修改类别和存放路径

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

sets = ['train', 'val', 'test']
classes = ["cat", "car", "dog", "person"]   # 改成自己的类别
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('Annotations存放路径/%s.xml' % (image_id), encoding='UTF-8')
    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
        #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('labels存放路径/'):
        os.makedirs('labels存放路径/')
    image_ids = open('ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
   
    if not os.path.exists('dataSet_path需要的存放路径/'):
        os.makedirs('dataSet_path需要的存放路径/')
     
    list_file = open('dataSet_path/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write('images存放路径/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

运行后会生成 labels 文件夹和 dataSet_path 文件夹，labels - 不同图像的标注文件

dataSet_path文件夹包含三个数据集的txt文件，train.txt等txt文件为划分后图像所在位置的路径

5. 配置文件

5.1 myvoc.yaml

在 yolov5 目录下的 data 文件夹下新建一个 myvoc.yaml文件，内容包括训练集以及验证集（train.txt和val.txt）的路径，目标的类别数目和类别名称

# :后需要空格
train:  # train.txt存放路径，不加引号
val:  # val.txt存放路径，不加引号

# number of classes
nc: 4

# class names
names: ["cat", "car", "dog", "person"]

如果yolov5/utils下有 autoanchor.py文件，就可以采用自动获取anchors。采用自动法，不用运行，训练时自动调用

5.2 修改模型配置文件

选择一个模型，在yolov5目录下的model文件夹下是模型的配置文件，有n、s、m、l、x版本，逐渐增大（随着架构的增大，训练时间也是逐渐增大）

官方数据：

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

yolov5/models:

5.3 修改参数

• 采用自动法获取anchors，只需更改nc 标注类别数，不用更改anchors

6. 模型训练

6.1 开始训练

python train.py --weights weights/yolov5s.pt  --cfg models/yolov5s.yaml  --data data/myvoc.yaml --epoch 200 --batch-size 8 --img 640   --device cpu

–epoch 200 ：训练200次

–batch-size 8：训练8张图片后进行权重更新

–device cpu：使用CPU训练。

6.2 训练可视化

tensorboard --logdir=runs

7. 效果检测

使用刚刚训练出的最好的模型 best.pt 来测试，在yolov5目录下的 runs/train/exp/weights ，测试结果保存在 yolov5/runs/detect 目录下

 python detect.py --weights runs/train/exp/weights/best.pt --source ../data/video/tram.mp4