Use Yolov5 to train your own data set, get started quickly

It summarizes the method of quickly getting started with Yolov5 to train the data set made by yourself . The steps are very detailed, and scholars are patient.

1. Prepare the Yolov5 framework

I have provided a source code package that has been debugged. The following tutorials are also explained based on the source code package provided by myself. Learners can download it by themselves. The link to download the source code package is: add link description, and the extraction code is : dq9k

Of course, you can also download the most original Yolov5 framework given by the official, and download it on Github . The download link is: Add link description
After opening the link, see the sample paper below. Download 5.0, which corresponds to the yolov5 framework:

2. Questions about datasets

The data set format used for Yolov5 training is yolo format, and the data set format used for Yolov3 training is VOC format. This is where the two training models use different data sets. Regarding the choice of data set format, I recommend the VOC format. It is directly made into the VOC data set format, and a piece of code can be used to directly convert the VOC data set into a yolo data set in the later stage. .

For details on the production of the VOC data set, please refer to my other blog, in which I introduced the production of the VOC data set in detail, the link is: Add link description

3. Convert VOC format data set to yolo format data set

Convert VOC format data set to yolo format data set, or use Yolov3 to store the data set file, copy the file containing VOC data set directly to the root directory of Yolov5, run the following code to directly convert the VOC format data set to yolo format The data set, the code will automatically separate the training set, the verification set, and the allocation ratio can be customized. How to copy files? See below:

The contents contained in the file VOCdevkit are as follows:

Run the following code directly in the Yolov5 project to convert the VOC format data set to the yolo format data set. The code is as follows:

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 = ["Before a whole","After a whole","Chest former","Chest after","Raise hand before","Raise hand after","Global left position","Global right position","Front face","Left face","Right face"]        ##这里要写好标签对应的类
# classes=["ball"]

TRAIN_RATIO = 80     #表示将数据集划分为训练集和验证集，按照2:8比例来的


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)
    out_file = open('VOCdevkit/VOC2007/YOLOLabels/%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/")
if not os.path.isdir(data_base_dir):
    os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
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)
yolov5_images_dir = os.path.join(data_base_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(data_base_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(wd, "yolov5_train.txt"), 'w')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
train_file.close()
test_file.close()
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
test_file = open(os.path.join(wd, "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:  # 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)
train_file.close()
test_file.close()

After running the above code, the newly converted yolo data set is shown below:

In the figure above, I have indicated which file stores the training data set, the verification set and which file stores the corresponding file location of the label.

4. Training model

After the data set is made, we can start training our own model. Here I provide a source code package that I have already debugged. It is recommended that learners download and debug according to the steps below. The probability of successful training is also higher. I It is also running after improvement on the basis of the official original framework. There are also many pits in the middle. These pits have been filled in my code, and learners can download and use them with confidence. In addition, my source code package also provides data sets. Learning can be used directly. The link to download the source package is: add link description , the extraction code is: dq9k

The folder sample paper after downloading and decompression is shown below:

Before officially starting the training, some modifications are required, see below:

(1) Modify the parameters of the configuration file in data:

(2) Modify the configuration file in the models file:

(3) After the configuration file is modified, modify the parameters in the train.py file:

After modifying the above parameters, you can directly run the train.py file to start training.

5. Enable tensorboard to view the parameter changes in the training process

Enter the following command in the terminal of the control panel. After running, a URL will be generated. Open this URL to view the changes in the training process in real time, see below:

tensorboard --logdir=runs

If the port to open the URL is occupied, you can also customize the port number by entering the following command:

tensorboard --logdir=runs --port==6007

The sample paper after opening the generated URL is shown below. From the figure, you can see the changes in each round of training process:

6. Problems that may be encountered during training

Problem 1: The virtual memory is not enough , the solution is as follows:

Problem 2: The video memory has exploded (overflowed) . The solution is to modify the batch-size. If the computer configuration is average, change it to a smaller size. See below:

Or you can solve it by the following method, modify the number of cores of the CPU , if the configuration is high, you can change it to a higher level, and the training will be faster, if the configuration is low, you can change it to a lower value so that there will be no overflow:

7. Detect the trained model

After training, use the trained model to detect pictures or videos.

(1) Image detection:

Modify some parameters in the file detect.py, see below:

Run the above modified detect.py file. After the detection is completed, the detection results will be stored in the runs–>detect–>exp path. Go to this path to view the detection results of the trained model detection, as follows:

(2) Video detection:
The parameters that need to be modified are as follows:

The detection results will also be stored under the path runs–>detect–>exp, just follow the prompts to view.

(3) Supplement: An error may be reported when calling the built-in camera of the computer, see below:

The solution is to modify the 280 lines of code in the datasets.py file in utils, see below:

Run the code after the modification, and you can call the computer’s built-in camera to detect normally, as follows:

The above is a quick way to use Yolov5 to train the data set made by yourself. Prepare the data set in advance. You only need to modify a few file parameters in the source package to train and use it. If the tutorial I summarized is helpful to you, thank you Support, thank you!