yolov5 (pytorch version) trains its own VOC format data set

I. Introduction

The previous two articles have explained the training process of the darknetversion yolov3and it yolov4's time to turn yolov5. I personally like it yolov5better, because it provides four different sizes of model structure to choose from, as shown in the figure below:

yolov5The training process is relatively simple, and it is also convenient windowsto proceed. After all, the pytorchframework is used. Briefly talk about how to train! ! !

2. Obtain and organize data sets

1. If you can download the labeled data set, you can use it directly. Pay attention to check whether the organization format of the data set is the same as the VOC format.

2. If you don't have a ready-made data set, you may need to label it yourself. I will elaborate on this later if I have the opportunity. Need to use some labeling tools, such as labelImg, the specific use method can refer to the following two articles! ! !

Use labelImg to label images under windows

labelImg installation instructions

3. The organization format of the data set

According to the official training tutorial, we need to convert the data set into a yoloformat, then we need to use a script to help us complete the conversion, but before that, you need to ensure that your data set organization structure is vocin the format. Regarding vocthe organization structure of the data set, please move to my other article: Detailed explanation of the organization structure of the VOC data set ! ! !

As shown in the figure below, this is the data set organization structure in the official training tutorial, and we need to be the same in the end, but the structure is the same, and the name does not have to be exactly the same. We only need to ensure that there images、labelsare folders and there are two subfolders under these two folders train、val. imagesThe image is labelsstored, and yolothe label of the format is stored .

note: images、labelstrain/valThe names of the subfolders under these two folders need to be consistent.

4. Format conversion

Here, you can move to another article of mine: VOC format data set to yolo (darknet) format . You need to understand the content of this article, otherwise you may not understand some of the things mentioned below, emmm! ! !

I converted the training verification data set and the test data set separately, so that it is convenient to save them separately, and then put them into the designated folders. The specific operations are as follows:

• The generated 2007_train.txtand 2007_val.txtmerged into one 2007_train.txt, because the training and verification pictures are in one piece, it is not easy to separate, so simply use them as training data, then we need to merge the two files into one, that is, directly 2007_val.txtcopy and paste the content into the 2007_train.txtfile , And then delete the 2007_val.txtfile.
• Then we use the test data set as the verification data, so we will 2007_test.txtmodify it to 2007_val.txt. Finally 2007_train.txt, 2007_val.txtput the files into imagesthe lower parallel path, this is not necessary, as long as you remember the path of these two files, and then later set up voc.yamldesignated to the correct time.
• By the above operation, we got 2007_train.txt, 2007_val.txtthe two documents, but also need to make some changes. Because the content of these two files is the full path of the training and verification pictures, but when we organize the data set later, the path of the pictures has changed. At this time, the pictures are all in the imagesfolder, so you need to combine these two The path before the picture name in each file is all replaced with the path actually located at this time. This can be done quickly by batch replacement.
• Put the labelscontents of the label storage folder generated by the two conversions into labels/trainand labels/valbelow respectively .
• Move JPEGImagesthe pictures in the respective folders of the training data set and the validation data set to images/trainand images/valbelow , respectively .

In this way, the organization of the data set is complete. However, in fact, you don’t necessarily need to do what I said above, because the point is that you need to know what is stored under each folder, and then convert it according to your own (even if you use both voc_label.pyto convert, but the operation may be different) The various files generated will finally imagesstore pictures and labelsstore yoloformat labels, and the training verification pictures and training verification labels can be stored in the trainand valsubfolders respectively .

For example, after you understand the following code, you can directly modify voc_label.pythe code in it to generate the required files more quickly. voc_label.pyThe content is as follows:

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
#转换最初的训练验证数据集的时候，使用以下sets
#转换最初的测试数据集的时候，可以将其改为sets=[('2007', 'test')]
sets=[('2007', 'train'), ('2007', 'val')]

classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", 
"car", "cat", "chair", "cow", "diningtable", "dog", "horse", 
"motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]


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(year, image_id):
    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
    out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, 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()

for year, image_set in sets:
    if not os.path.exists('VOCdevkit/VOC%s/labels/'%(year)):
        os.makedirs('VOCdevkit/VOC%s/labels/'%(year))
    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'
    							%(year, image_set)).read().strip().split()
    
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    
    for image_id in image_ids:
        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n'
        						%(wd, year, image_id))
        convert_annotation(year, image_id)
    list_file.close()

Three, build the project

1. Download the source code

https://github.com/ultralytics/yolov5

2. Environment configuration

pip install -U -r requirements.txt

Fourth, modify relevant documents

1.data/voc.yaml File

As shown below, we need to train、valset the value as said before 2007_train.txt, 2007_val.txtthe storage path of the file. nc和namesYou also need to modify it according to your actual situation, which respectively indicate the number of categories and the names of the categories.

# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
# Train command: python train.py --data voc.yaml
# Default dataset location is next to /yolov5:
#     /parent_folder
#     /VOC
#     /yolov5


# download command/URL (optional)
download: bash data/scripts/get_voc.sh

# train and val data as 
# 1) directory: path/images/, 
# 2) file: path/images.txt, or 
# 3) list: [path1/images/, path2/images/]
train: F:\0-My_projects\My_AI_Prj\0_PyTorch_projects\ultralytics_yolov5s_m_l_x_pytorch
\my_dataset\2007_train.txt  # 16551 images

val: F:\0-My_projects\My_AI_Prj\0_PyTorch_projects\ultralytics_yolov5s_m_l_x_pytorch
\my_dataset\2007_val.txt  # 4952 images

# number of classes
nc: 20

# class names
names: ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 
'cat', 'chair', 'cow', 'diningtable', 'dog','horse', 'motorbike', 
'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']

2.models/yolov5s.yaml File

Here I am preparing for training yolov5s. If you want to train the other three models, just select the corresponding configuration file. We generally only need to modify the ncnumber of types in the first row for our own data set.

3. Training parameters

For convenience, we can train.pymodify the parameters directly in the file, and then run it directly train.py, as shown below:

It should be noted that if weightsthe path you specify , that is, the first parameter, if it does not exist, then it will download it by itself. If the download fails, then you can download the initial weight file by yourself, and then put it under this path.

Five, start training

Just run the train.pyfile directly .