利用yolo5官方模型训练自己的数据集_yolo官方

利用yolov5官方模型权重训练自己的数据集

一、运行环境
Cuda = 10.1
python=3.7.x
torch =1.6.0
cython = 0.29.21
numpy = 1.19.1
opencv-python = 4.3.0.38
matplotlib = 3.3.0
pillow = 7.2.0
tensorboard = 2.2.0
pyYAML = 5.3.1
torchvision = 0.7.0
scipy = 1.5.2
tqdm = 4.48.2
二、yolov5文件
数据集文件：coco128
权重文件
源码文件
以上数据在百度云
链接：https://pan.baidu.com/s/1ReSfDQojgBFg6XFOjAn8og
提取码：6tzx
三、使用自己的数据集训练数据
1、仿照coco128.yaml文件写自己的data.yaml文件，在coco128.yaml文件的文件内容如下：

# COCO 2017 dataset http://cocodataset.org - first 128 training images
# Train command: python train.py --data coco128.yaml
# Default dataset location is next to /yolov5:
#   /parent_folder
#     /coco128
#     /yolov5


# download command/URL (optional)
download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip

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

# number of classes
nc: 80

# class names
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']
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

其中需要自己仿照的部分：
train：自己数据集的相对路径
val：自己的val文件相对路径
nc：自己定义的类别的数量（其实可以不需要修改前提是自己的类别小于80）
names：自己定义类别的名称（如果上面nc没有修改只需要修改自己类别的前几个即可）
2.标签文件建立
可以使用labelImg标注图像，具体使用步骤可查看https://blog.csdn.net/python_pycharm/article/details/85338801
该标注软件标注之后是xml文件需要对xml文件进行处理，处理成yolo能识别的txt文件。
具体代码如下：
（1）获取所有训练集图片的文件目录

import os

file_path = "F:/code/20200414/region1" # 训练集图片所属文件夹
path_list = os.listdir(file_path)  # os.listdir(file)会历遍文件夹内的文件并返回一个列表
print(path_list)
path_name = []  # 把文件列表写入save.txt中


def saveList(pathName):
    for file_name in pathName:
        with open("F:/code/20200414/train.txt", "a") as f:
            f.write(file_name.split(".")[0] + "\n")
            


def dirList(path_list):
    # 将文件路径从列表中获取出来
    for i in range(0, len(path_list)):
        path = os.path.join(file_path, path_list[i])

    if os.path.isdir(path):
        print(path)
        saveList(os.listdir(path))
        


if __name__ == '__main__':
    dirList(path_list)
    saveList(path_list)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

（2）将xml文件转换为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 = ['1'] # 注意这个部分需要修改成属于自己自定义的类别
# 进行归一化操作
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]

def convert_annotation(image_id):
    print('*'*100)
    '''
    将对应文件名的xml文件转化为label文件，xml文件包含了对应的bunding框以及图片长款大小等信息，
    通过对其解析，然后进行归一化最终读到label文件中去，也就是说
    一张图片文件对应一个xml文件，然后通过解析和归一化，能够将对应的信息保存到唯一一个label文件中去
    labal文件中的格式：calss x y w h　　同时，一张图片对应的类别有多个，所以对应的ｂｕｎｄｉｎｇ的信息也有多个
    '''
    # 对应的通过year 找到相应的文件夹，并且打开相应image_id的xml文件，其对应bund文件
    in_file = open('C:/Users/workAI/Desktop/data/Annotations/%s.xml' % image_id)
    # print(in_file.name)
    # 准备在对应的image_id 中写入对应的label，分别为
    # <object-class> <x> <y> <width> <height>
    out_file = open('C:/Users/workAI/Desktop/data/labels/%s.txt' % image_id, 'w')
    # print(out_file.name)
    # 解析xml文件
    tree = ET.parse(in_file)
    # 获得对应的键值对
    root = tree.getroot()
    # 获得图片的尺寸大小
    size = root.find('size')
    # 获得宽
    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)
        print(bb)
        # 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('C:/Users/workAI/Desktop/data/labels/'):
        os.makedirs('C:/Users/workAI/Desktop/data/labels/')
    # 读取在ImageSets/Main 中的train、test..等文件的内容
    # 包含对应的文件名称
    image_ids = open('C:/Users/workAI/Desktop/data/ImageSets/%s.txt' % image_set).read().strip().split()
    # 打开对应的2012_train.txt 文件对其进行写入准备
    list_file = open('C:/Users/workAI/Desktop/data/%s.txt' % image_set, 'w')
    # 将对应的文件_id以及全路径写进去并换行
    for image_id in image_ids:
        list_file.write('C:/Users/workAI/Desktop/data/images/%s.jpg\n' % image_id)
        # 调用  year = 年份  image_id = 对应的文件名_id
        convert_annotation(image_id)
    # 关闭文件
    list_file.close()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

3.文件组织目录
将数据集coco128（换成自己的数据集）放在yolov5的源码旁边。数据集coco128文件夹中包含images和labels两个文件夹，分别存放数据集的图片和数据集的标注信息。
4.模型选择
在models文件中官方给定了五个模型，自己只需要在模型中更改自己定义的nc的数量大小即可。
四、训练
1、train.py文件
参数：
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：单类别的训练集
2、在训练的结束时会产生runs文件夹，在文件夹中会有

在这个训练轮次中会有

在weights中会有

五、训练模型测试
detect.py文件
将上述文件的weights文件夹下的best.pt文件移动到yolov5文件夹下，并在detect.py文件修改

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='best.pt', help='model.pt path(s)')
1
2
3

即可测试，
在测试时需要将自己测试的图片放到yolov5\inference\images文件夹下，在测试过程中会产生yolov5\inference\output文件夹，并在文件夹中会产生标注的文件。