YOLOv5训练自己的数据集

一、概述

  使用YOLOv5训练数据集，大致分为3个步骤：训练前的数据准备及处理、训练自己的数据集、检测，以及训练后一些优化问题。

二、训练前准备工作

首先从YOLOv5的官方网站：https://github.com/ultralytics/yolov5 下载对应项目到自己的平台，我这里使用的是Google cloab.

1.处理数据

  对你的图片用labelImge之类的软件进行处理，标记你的图片，并且注意分别保存xml文件和图片

 如果需要对你的图片重新进行批量命名，要命名后再标注图片

#批量重命名文件
import os
import sys
def re_name():
    path=input("请输入路径(例如D:\\\\picture)：")
    name=input("请输入开头名:")
    startNumber=input("请输入开始数:")
    fileType=input("请输入后缀名（如 .jpg、.txt等等）:")
    print("正在生成以"+name+startNumber+fileType+"迭代的文件名")
    count=0
    filelist=os.listdir(path)
    for files in filelist:
        Olddir=os.path.join(path,files)
        if os.path.isdir(Olddir):
            continue
        Newdir=os.path.join(path,name+str(count+int(startNumber))+fileType)
        os.rename(Olddir,Newdir)
        count+=1
    print("一共修改了"+str(count)+"个文件")

re_name()

2.环境配置

（1）直接通过pip install -r requirements.txt 直接下载,该文件在下载好的YOLOv5项目文件中，要注意路径问题。

（2）通过手动pip下载

pip install numpy
pip install matplotlib
pip install pandas
pip install scipy
pip install seaborn
pip install opencv-python
pip install tqdm
pip install pillow
pip install tensorboard
pip install pyyaml
pip install pandas
pip install scikit-image
pip install Cython
pip install thop
pip install pycocotools

 3.建立对应的文件夹

  

 文件最好在对应的项目中，运行时会调用项目中对应的模块。其中：

Annotations：存放xml形式的标签文件

ImageSets:存放的是分类和检测的数据集分割文件，包含train.txt（包含用于训练的图片名称）, val.txt（写着用于验证的图片名称）,trainval.txt（train与val的合集）,test.txt（写着用于测试的图片名称），这一部分后期代码会自动生动，可以不用建立

JPEGImages:存放对应img图片

labels：存放label标注信息的txt文件，并与图片相对应

4.处理数据

将标注过的图片，以及xml文件分别传入JPEGImages、Annotations中，然后我们来建立一个makeTxt.py文件来分配用于训练和验证的数据。

import os
import random


trainval_percent = 0.9
train_percent = 0.9  #用于训练的数据的比例，不能设置为1，会导致
xmlfilepath = 'data/Annotations'
txtsavepath = 'data/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)
train = random.sample(trainval, tr)

ftrainval = open('data/ImageSets/trainval.txt', 'w')
ftest = open('data/ImageSets/test.txt', 'w')
ftrain = open('data/ImageSets/train.txt', 'w')
fval = open('data/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:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)

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

运行完后会在ImageSets中产生以下文件

然后再建立voc_label.py,在这里要把classes=[  ]中的标签改为自己的具体名称，运行完后会在labels文件中产生对应标签的txt文件。

# xml解析包
import xml.etree.ElementTree as ET
import pickle
import os
# os.listdir() 方法用于返回指定的文件夹包含的文件或文件夹的名字的列表
from os import listdir, getcwd
from os.path import join


sets = ['train', 'test', 'val']
classes = ['computer', 'person','phone','tablet phone','cup','bag','bag2','books']


# 进行归一化操作
def convert(size, box): # size:(原图w,原图h) , box:(xmin,xmax,ymin,ymax)
    dw = 1./size[0]     # 1/w
    dh = 1./size[1]     # 1/h
    x = (box[0] + box[1])/2.0   # 物体在图中的中心点x坐标
    y = (box[2] + box[3])/2.0   # 物体在图中的中心点y坐标
    w = box[1] - box[0]         # 物体实际像素宽度
    h = box[3] - box[2]         # 物体实际像素高度
    x = x*dw    # 物体中心点x的坐标比(相当于 x/原图w)
    w = w*dw    # 物体宽度的宽度比(相当于 w/原图w)
    y = y*dh    # 物体中心点y的坐标比(相当于 y/原图h)
    h = h*dh    # 物体宽度的宽度比(相当于 h/原图h)
    return (x, y, w, h)    # 返回 相对于原图的物体中心点的x坐标比,y坐标比,宽度比,高度比,取值范围[0-1]


# year ='2012', 对应图片的id（文件名）
def convert_annotation(image_id):
    '''
    将对应文件名的xml文件转化为label文件，xml文件包含了对应的bunding框以及图片长款大小等信息，
    通过对其解析，然后进行归一化最终读到label文件中去，也就是说
    一张图片文件对应一个xml文件，然后通过解析和归一化，能够将对应的信息保存到唯一一个label文件中去
    labal文件中的格式：calss x y w h　　同时，一张图片对应的类别有多个，所以对应的ｂｕｎｄｉｎｇ的信息也有多个
    '''
    # 对应的通过year 找到相应的文件夹，并且打开相应image_id的xml文件，其对应bund文件
    in_file = open('data/Annotations/%s.xml' % (image_id), encoding='utf-8')
    # 准备在对应的image_id 中写入对应的label，分别为
    #     
    out_file = open('data/labels/%s.txt' % (image_id), 'w', encoding='utf-8')
    # 解析xml文件
    tree = ET.parse(in_file)
    # 获得对应的键值对
    root = tree.getroot()
    # 获得图片的尺寸大小
    size = root.find('size')
    # 如果xml内的标记为空，增加判断条件
    if size != None:
        # 获得宽
        w = int(size.find('width').text)
        # 获得高
        h = int(size.find('height').text)
        # 遍历目标obj
        for obj in root.iter('object'):
            # 获得difficult ？？
            difficult = obj.find('difficult').text
            # 获得类别 =string 类型
            cls = obj.find('name').text
            # 如果类别不是对应在我们预定好的class文件中，或difficult==1则跳过
            if cls not in classes or int(difficult) == 1:
                continue
            # 通过类别名称找到id
            cls_id = classes.index(cls)
            # 找到bndbox 对象
            xmlbox = obj.find('bndbox')
            # 获取对应的bndbox的数组 = ['xmin','xmax','ymin','ymax']
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            print(image_id, cls, b)
            # 带入进行归一化操作
            # w = 宽, h = 高， b= bndbox的数组 = ['xmin','xmax','ymin','ymax']
            bb = convert((w, h), b)
            # bb 对应的是归一化后的(x,y,w,h)
            # 生成 calss x y w h 在label文件中
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


# 返回当前工作目录
wd = getcwd()
print(wd)


for image_set in sets:
    '''
    对所有的文件数据集进行遍历
    做了两个工作：
　　　　１．将所有图片文件都遍历一遍，并且将其所有的全路径都写在对应的txt文件中去，方便定位
　　　　２．同时对所有的图片文件进行解析和转化，将其对应的bundingbox 以及类别的信息全部解析写到label 文件中去
    　　　　　最后再通过直接读取文件，就能找到对应的label 信息
    '''
    # 先找labels文件夹如果不存在则创建
    if not os.path.exists('data/labels/'):
        os.makedirs('data/labels/')
    # 读取在ImageSets/Main 中的train、test..等文件的内容
    # 包含对应的文件名称
    image_ids = open('data/ImageSets/%s.txt' % (image_set)).read().strip().split()
    # 打开对应的2012_train.txt 文件对其进行写入准备
    list_file = open('data/%s.txt' % (image_set), 'w')
    # 将对应的文件_id以及全路径写进去并换行
    for image_id in image_ids:
        list_file.write('data/images/%s.jpg\n' % (image_id))
        # 调用  year = 年份  image_id = 对应的文件名_id
        convert_annotation(image_id)
    # 关闭文件
    list_file.close()

5.部分文件参数的修改

（1）data文件下数据集文件coco.yaml的修改或者可以复制代码重新建一个文件，将数据路径（train、val、test）修改成前面ImageSet下对应的txt文件，将nc(数据中标签个数)，以及names中的标签改为自己的

# YOLOv5  by Ultralytics, GPL-3.0 license
# COCO 2017 dataset http://cocodataset.org
# Example usage: python train.py --data coco.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco  ← downloads here


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco          # dataset root dir
train:   #改为自己对应的txt文件  # train images (relative to 'path') 118287 images
val:                            # train images (relative to 'path') 5000 images
test:                           # 20288 of 40670 images, submit to 

# Classes
nc: 80  # number of classes
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']  # class names

(2)使用模型的修改，使用yolov5s.pt,就修改对应的yolov5-master/models/yolov5m.yaml，将nc修改为自己对应的标签个数，这里还涉及到一个anchors box的设定问题.

 (3)train.py的修改

  主要是--weight、--cfg、--data、--hyp、--epochs、--batch-size、--imagsz几个参数，--cache容易报错，最好也修改一下。

#在default中写入自己数据的路径
#--weight下载对应的yolov5s.pt、yolov5m.pt等到自己的项目中，使用哪个就填写对应的路径
parser.add_argument('--weights', type=str, default='', help='initial weights path')
#--cfg对应你的--weight，在models中，对应前面5.（2）的文件
parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
 #对应前面5.（1）的文件
parser.add_argument('--data', type=str, default='', help='dataset.yaml path')
parser.add_argument('--hyp', type=str, default='data/hyps/hyp.scratch.yaml', help='hyperparameters path')
#迭代次数
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
#图片尺寸
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
parser.add_argument('--cache', action='store_true', help='cache images for faster training')

 三、数据训练

在完成以上步骤后，就可以运行train.py文件对自己的数据开始训练，在出现下面的界面时，就显示训练已经正常开始进行。

 四、检测

    对detect.py文件进行修改，主要是--weight训练好的权重的选择、--source检测图片的路径，然后运行py文件就可以在相应文件夹中得到对应的结果。

#选择权重
parser.add_argument('--weights', nargs='+', type=str, default='选择自己训练好的权重', help='model.pt path(s)')
#选择检测图片路径
parser.add_argument('--source', type=str, default='选择需要检测图片的路径', help='file/dir/URL/glob, 0 for webcam')
#运行结果保存的路径
parser.add_argument('--project', default='runs/detect', help='save results to project/name')

附带一张我自己的检测图

 五、优化（识别率的提高）

1.数据集的数量：增大数据集的数量可以明显提高识别的准确度

2.识别物体的特征值：识别的物体是一类物体对应一个标签，具有较一致的特征值，尽量细分，可以提高识别准确度

3.修改anchor box,项目中的anchor值可能并不适合我们自己的，我们可以通过k-means计算自己的anchor box(目前还在研究)，anchor值应该是会随着我们自己的数据发生改变的，但是实际情况却没有，可以尝试更改yolov5-master/utils/autoanchor.py里面的kmean_anchors函数的参数，再跑一下代码

def kmean_anchors(dataset='写入自己的数据路径.yaml', n=8, img_size=640, thr=4.0, gen=1000, verbose=True):

4.更改模型，我尝试了yolov5s.pt以及yolov5m.pt，m比s效果更好，不过m产生的权重会更大，大约在50MB，s为14MB.