YOLOv5 从0到1训练自己数据

提示

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

文章目录

• 一、配置环境
• 二、使用步骤
  • 1.引入库
  • 2.读入数据

一、环境配置

环境：ubuntu16。

1.1 下载源码： 

官网源码下载：  https://github.com/ultralytics/yolov5

1.2创建虚拟环境：

conda create -n yolov5 python==3.8
conda activate yolov5
pip install -r requirements.txt
或者
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt

二、制作数据集

2.1.labelimg标注

2.1.1 安装环境

安装环境如下： win10+python3.6+PyQt5+lxml+labelImg。其中Python3.6和PyQt5、lxml为Anaconda下自动安装，因此下载安装Anaconda后的朋友可以直接忽略环境的安装。

从github下载源码并解压，得到文件夹LabelImg-master，github地址：https://github.com/tzutalin/labelImg

输入：pyrcc5 -o resources.py resources.qrc

运行成功后再次输入：python labelImg.py

 遇到该错误提示后，将resources.py文件复制进入libs文件夹中。

cmd重新运行：python labelImg.py

创建类别classes.txt，放置于和图片同一个文件夹下 ：

2.2.数据结构

首先建立一个自己的数据文件夹：mycoco

其目录结构如下：（剩下的文件架构代码生成）

mycoco
├── all_images
├── all_xml
├── make_txt.py
└── train_val.py

其中all_images文件夹下放置所有图片，all_xml文件夹下放置所有与之对应的xml文件。

make_txt.py文件是用来划分数据集使用，内容如下：

import os
import random
trainval_percent = 0.1  
train_percent = 0.9    
xmlfilepath = 'all_images'
txtsavepath = 'ImageSets'
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) #从所有list中返回tv个数量的项目
train = random.sample(trainval, tr)
if not os.path.exists('ImageSets/'):
    os.makedirs('ImageSets/')
ftrainval = open('ImageSets/trainval.txt', 'w')
ftest = open('ImageSets/test.txt', 'w')
ftrain = open('ImageSets/train.txt', 'w')
fval = open('ImageSets/val.txt', 'w')
for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

运行python make_txt.py 结果，生成四个只包含图片名称的txt文件：

 接着运行train_val.py,该文件一方面将all_xml中xml文件转为txt文件存于all_labels文件夹中，另一方面生成训练所需数据存放架构。（这里如果你的数据直接是txt的标签的话将标签转化的功能注释掉即可）代码如下：

import xml.etree.ElementTree as ET
import pickle
import os
import shutil
from os import listdir, getcwd
from os.path import join
sets = ['train', 'trainval']
classes = ['car','chemicals vehicle','truck','bus','triangle warning sign','warning sign','warning slogan']
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    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('all_xml/%s.xml' % (image_id))
    out_file = open('all_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()
print(wd)
for image_set in sets:
    if not os.path.exists('all_labels/'):
        os.makedirs('all_labels/')
    image_ids = open('ImageSets/%s.txt' % (image_set)).read().strip().split()
    image_list_file = open('images_%s.txt' % (image_set), 'w')
    labels_list_file=open('labels_%s.txt'%(image_set),'w')
    for image_id in image_ids:
        image_list_file.write('%s.jpg\n' % (image_id))
        labels_list_file.write('%s.txt\n'%(image_id))
        convert_annotation(image_id) #如果标签已经是txt格式，将此行注释掉，所有的txt存放到all_labels文件夹。
    image_list_file.close()
    labels_list_file.close()


def copy_file(new_path,path_txt,search_path):#参数1：存放新文件的位置  参数2：为上一步建立好的train,val训练数据的路径txt文件  参数3：为搜索的文件位置
    if not os.path.exists(new_path):
        os.makedirs(new_path)
    with open(path_txt, 'r') as lines:
        filenames_to_copy = set(line.rstrip() for line in lines)
        # print('filenames_to_copy:',filenames_to_copy)
        # print(len(filenames_to_copy))
    for root, _, filenames in os.walk(search_path):
        # print('root',root)
        # print(_)
        # print(filenames)
        for filename in filenames:
            if filename in filenames_to_copy:
                shutil.copy(os.path.join(root, filename), new_path)

#按照划分好的训练文件的路径搜索目标，并将其复制到yolo格式下的新路径
copy_file('./images/train/','./images_train.txt','./all_images')
copy_file('./images/val/','./images_trainval.txt','./all_images')
copy_file('./labels/train/','./labels_train.txt','./all_labels')
copy_file('./labels/val/','./labels_trainval.txt','./all_labels')

运行python train_val.py，结果

mytrain
├── mycoco
│   ├── all_images
│   ├── all_labels
│   ├── all_xml
│   ├── ImageSets
│   │   ├── train.txt
│   │   ├── test.txt
│   │   ├── trainval.txt
│   │   └── val.txt
│   ├── images
│   │   ├── train
│   │   └── val
│   ├── labels
│        ├── train
│        └── val
│   ├── images_train.txt
│   ├── images_trainval.txt
│   ├── labels_train.txt
│   ├── labels_trainval.txt
│   ├── make_txt.py
│   └── train_val.py
└── yolov5

 至此数据集已经完全建好。将制作好的mycoco文件夹与下载好的yolov5文件夹放入同一级文件夹中。

接着按照yolov5-master/data/coco128.yaml文件，制作mycoco.yaml文件(与coco128.yaml文件同目录)：

# Default dataset location is next to /yolov5:
#   /parent_folder
#     /mycoco
#     /yolov5


# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../mycoco/images/train/  
val: ../mycoco/images/val/  

# number of classes
nc: 7

# class names
names: ['car','chemicals vehicle','truck','bus','triangle warning sign','warning sign','warning slogan' ]

三、训练

3.1.训练准备

1. 在YOLOv5的GitHub开源网址上下载对应版本的模型：Releases · ultralytics/yolov5 · GitHub

 

接着打开train.py文件，train.py文件详解，更改相应的参数（预训练模型，训练的数据集），这里使用yolov5s.pt为预训练模型，更改yolov5s.yaml文件中的参数（自己的类别数）

参数解释：

epochs：指的就是训练过程中整个数据集将被迭代多少次,显卡不行你就调小点。
batch-size：一次看完多少张图片才进行权重更新，梯度下降的mini-batch,显卡不行你就调小点。
cfg：存储模型结构的配置文件
data：存储训练、测试数据的文件
img-size：输入图片宽高,显卡不行你就调小点。

python train.py --img 640 --batch 32 --epoch 300 --data data/mydata.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --device '0,1' 

3.2.测试数据

3.2.数据增广：

YOLOv5采用了许多有效的数据处理的方法来增加训练模型的精度和减少训练时间，数据扩增的方法主要介绍Mosaic、Cutout方法，除了这两种方法外yolov5还使用了图像扰动，改变亮度、对比对、饱和度、色调，加噪声，随机缩放，随机裁剪（random crop），翻转，旋转，随机擦除等

Mosaic简单的说就是把四张训练图片缩放拼成一张图，丰富检测物体的背景，并且四张图片拼接在一起变相地提高了batch_size，Mosaic有利于提升小目标的检测，这是因为一般在数据集中小目标在图片中分布不均匀，这导致在常规的训练中小目标的学习总是不太充分。使用mosaic数据增强后，在遍历每个张图片包含了四张图片具有小目标的可能性就很大了，同时，每张图都有不同程度的缩小，即使没有小目标，通过缩小，原来的目标尺寸也更接近小目标的大小，这对模型学习小目标很有利。

Cutout就是随机选择一个固定大小的正方形区域，然后采用全0填充就OK了，当然为了避免填充0值对训练的影响，应该要对数据进行中心归一化操作，norm到0。

3.2.1 启用数据增强：

1.data目录下，有hyp文件：data/hyp.finetune.yaml:

在train.py中添加指定:

YOLOV5源码解读（数据集加载和增强）_暮丶凉的博客-CSDN博客\

yolov5数据增强代码解读_weixin_41012399的博客-CSDN博客_yolov5中的数据增强

《YOLOV4&5原理与源代码解析之一：数据增强》 - JackRuiYu - 博客园

yolov5代码解读-dataset - 知乎

yolov5代码阅读笔记_RainbowSun1102的博客-CSDN博客_yolov5代码

yolov5启用数据增强、tensorboard可视化及cutout马赛克增强_小俊俊的博客-CSDN博客

 4.问题锦集

参考：

YOLOv5训练自己的数据集(超详细)_AI小白一枚的博客-CSDN博客_yolov5训练自己的数据yolov5训练自己的数据集（一文搞定训练）_小屋-CSDN博客[图像算法]-YOLOV5测试及训练自己的数据集 - 简书

Yolov5-Pytorch版-Windows下训练自己的数据集，内含voc批量转yolo方法。（自称宇宙超级巨详细步骤）_136.la

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