Yolov5训练自己的数据集_yolov5数据集下载

【前段时间往硬盘里导东西，不小心丢了好多项目和数据，其中包括之前调试好的Yolov5和用小孔数据集训练好的权重文件、预测结果等等等，恢复无果，今天重新配置一番，发现之前的配置和训练过程没有记录，好多步骤都忘了，因此写下这篇文章做记录】

1. 下载

需要下载的内容就两个：Yolov5项目包 and 预训练权重文件，搭梯子下载快一点。
Yolov5项目包下载地址：https://github.com/ultralytics/yolov5/tree/v5.0
权重文件下载地址 ：https://github.com/ultralytics/yolov5/releases
权重文件我一般习惯放在weights这个文件夹里。

下载好后还需要按照requirements.txt文件配置环境，配置好环境后可以运行detect.py文件测试一下，解决解决报错。
这里记录两个报错，解决方法也都是参考的其他文章。

AttributeError: Can‘t get attribute ‘SPPF‘ on ＜module ‘models.common‘ from XXXX/model/common.py

原因是common.py中没有SPPF这个类，应该是版本问题。其他文章中给出两个解决方法，一是下载其他版本中的commom.py文件替换，二是在common.py文件中加上这个类，注意别忘了import warnings

import warnings

class SPPF(nn.Module):
    # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
    def __init__(self, c1, c2, k=5):  # equivalent to SPP(k=(5, 9, 13))
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * 4, c2, 1, 1)
        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
            y1 = self.m(x)
            y2 = self.m(y1)
            return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

RuntimeError: The size of tensor a (80) must match the size of tensor b (56) at non-singleton

原因是权重文件有问题，我直接用这个链接里的v5s权重文件替换了自己的s权重文件就解决了。

2. 数据集制作和转换

首先是制作自己的数据集，labelimg标注，不多说。
接着是将标注文件转换为Yolo可用的格式，我一般习惯标注为voc格式，再转换为yolo可用的txt格式。
代码直接用的炮哥的https://blog.csdn.net/didiaopao/article/details/120022845文章里给出了voc格式转换为txt格式并自动划分训练集和验证集的代码，还给出了txt格式直接划分训练集和验证集的代码。
下面粘的代码是我根据自己的文件路径修改过的，并且具有无缺陷图片生成空txt文件的功能（我的数据集正常图片是没有标注文件的，原代码中会直接略过这些图片，所以我修改了一下，粘在这里只是防止自己下次弄丢资料，方便来找）

#功能：将xml格式标注文件转换为txt标注文件，并自动划分训练集和验证集
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile

#classes = ["hat", "person"]
# classes=["ball"]
classes = ["cracks", "sandholes",'bumps']
TRAIN_RATIO = 80


def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)


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('VOCdevkit/VOC2007/Annotations/%s.xml' % image_id)
    xml_path = os.path.join(annotation_dir,image_id+'.xml')
    in_file = open(xml_path)
    #out_file = open('VOCdevkit/VOC2007/YOLOLabels/%s.txt' % image_id, 'w')
    out_file = open(yolo_labels_dir+'/%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')
    in_file.close()
    out_file.close()

wd = os.getcwd()
wd = os.getcwd()
#data_base_dir = os.path.join(wd, "VOCdevkit/")
data_base_dir  = 'E:/Dataset/'
dataset_name = 'Ostiole_Samples'
if not os.path.isdir(data_base_dir):
    os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, dataset_name+'/')
if not os.path.isdir(work_sapce_dir):
    os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
    os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
    os.mkdir(image_dir)
clear_hidden_files(image_dir)
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
if not os.path.isdir(yolo_labels_dir):
    os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)
save_dir = os.path.join(work_sapce_dir, "forYolo/")
if not os.path.isdir(save_dir):
    os.mkdir(save_dir)
clear_hidden_files(save_dir)
yolov5_images_dir = os.path.join(save_dir, "images/")
if not os.path.isdir(yolov5_images_dir):
    os.mkdir(yolov5_images_dir)
clear_hidden_files(yolov5_images_dir)
yolov5_labels_dir = os.path.join(save_dir, "labels/")
if not os.path.isdir(yolov5_labels_dir):
    os.mkdir(yolov5_labels_dir)
clear_hidden_files(yolov5_labels_dir)
yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
if not os.path.isdir(yolov5_images_train_dir):
    os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
if not os.path.isdir(yolov5_images_test_dir):
    os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
if not os.path.isdir(yolov5_labels_train_dir):
    os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)
yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
if not os.path.isdir(yolov5_labels_test_dir):
    os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)

train_file = open(os.path.join(save_dir, "yolov5_train.txt"), 'w')
test_file = open(os.path.join(save_dir, "yolov5_val.txt"), 'w')
train_file.close()
test_file.close()
train_file = open(os.path.join(save_dir, "yolov5_train.txt"), 'a')
test_file = open(os.path.join(save_dir, "yolov5_val.txt"), 'a')
list_imgs = os.listdir(image_dir)  # list image files
prob = random.randint(1, 100)
print("Probability: %d" % prob)
for i in range(0, len(list_imgs)):
    path = os.path.join(image_dir, list_imgs[i])
    if os.path.isfile(path):
        image_path = image_dir + list_imgs[i]
        voc_path = list_imgs[i]
        (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
        (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
        annotation_name = nameWithoutExtention + '.xml'
        annotation_path = os.path.join(annotation_dir, annotation_name)
        label_name = nameWithoutExtention + '.txt'
        label_path = os.path.join(yolo_labels_dir, label_name)
    prob = random.randint(1, 100)
    print("Probability: %d" % prob)
    if (prob < TRAIN_RATIO):  # train dataset
        if os.path.exists(annotation_path):
            train_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_train_dir + voc_path)
            copyfile(label_path, yolov5_labels_train_dir + label_name)
        else:
            train_file.write(image_path + '\n')
            open(label_path,'w').close()
            copyfile(image_path, yolov5_images_train_dir + voc_path)
            copyfile(label_path, yolov5_labels_train_dir + label_name)
    else:  # test dataset
        if os.path.exists(annotation_path):
            test_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_test_dir + voc_path)
            copyfile(label_path, yolov5_labels_test_dir + label_name)
        else:

            test_file.write(image_path + '\n')
            open(label_path,'w').close()
            copyfile(image_path, yolov5_images_test_dir + voc_path)
            copyfile(label_path, yolov5_labels_test_dir + label_name)
train_file.close()
test_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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164

分完训练集和验证集之后就像下图这样。

YOLOLabels中存放的是所有图片的txt格式标注文件，没有标注的就是空文件。另外还生成两个txt文件，分别存放训练集和验证集图片的路径。

修改配置文件

环境配置好、数据集制作完成，就可以进行训练了。每次训练时都需要修改两个配置文件，一个是data里的.yaml文件，一个是models里的.yaml文件。

• data中的.yaml文件需要给出训练集和验证集路径、类别数量和类别名称，下面是我的data/.yaml文件

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

# number of classes
nc: 3

# class names
names: ['cracks','sandholes','bumps']
1
2
3
4
5
6
7
8
9

• models下的yaml文件存放的是模型的超参数，一般只需要修改第一行类别数就可以了。
  
  然后就可以开始训练了，我一般习惯在程序里直接修改参数进行训练，需要修改的一般是weights、cfg、data，可能还有epochs、batch-size，内存报错的时候可能还需要调小workers。
  
  这里记录一个报错，第一次调试的时候没有，但是这次出现了。

RuntimeError: result type Float can‘t be cast to the desired output type long int

解决方法参考的是这篇文章