[AI combat] Hands-on training of your own target detection model (YOLO)

In the previous article, we have introduced the use of your own data to train the target detection model based on SSD (see the article: Teach you to train your own target detection model ), this article will introduce how to use your own data based on another target detection model YOLO to train.

YOLO (You only look once) is one of the most popular target detection models. The latest version has been developed to V3, and it is widely used in the industry. The basic principle of YOLO is: first divide the input image into 7x7 grids, predict 2 borders for each grid, and then remove the target windows with low probability according to the threshold, and finally remove redundant windows by border merging. , and the detection result is obtained, as shown in the figure below:

YOLO is characterized by "fast", but because YOLO only predicts one object for each grid, it is easy to cause missed detection, and it is relatively sensitive to the scale of the object. Object generalization ability is poor.

The goal of this paper is still to identify and detect cute pandas in images

Based on YOLO, use its own data to train the target detection model. The training process is the same as the SSD-based training model introduced in the previous article. The main steps are as follows:

1. Install the labeling tool
The labeling tool used in this case is labelImg. In the previous article, the training SSD was introduced. The installation method is described in detail in the model (see the article: Teach you how to train your own target detection model ), so I won't repeat it here.
The labelImg labeling tool after successful installation, as shown below:

2. Labeling data
Use the labelImg tool to frame the panda photo, automatically generate an xml file in VOC_2007 format, and save it as a training data set. The operation method is the same as the labeling method for training the SSD model in the previous article (see the article: Teach you to train your own target detection model ), so I won't repeat it here.

3. Configure YOLO
(1) Install Keras

This case uses the latest V3 version of YOLO, based on the Keras version. Keras is a high-level neural network API with Tensorflow, Theano and CNTK as backends. Since the basic environment of this case (see the article: AI Basic Environment Construction ) has already installed tensorflow, the bottom layer of Keras will call tensorflow to run the model. Keras is installed as follows:

# 切换虚拟环境
source activate tensorflow
# 安装keras-gpu版本
conda install keras-gpu
# 如果是安装 keras cpu版本，则执行以下指令
#conda install keras

The keras version of yolo3 also depends on the PIL toolkit. If it has not been installed before, it should also be installed in anaconda

# 安装 PIL
conda install pillow

(2) Download yolo3 source code Download the source code
on the github of keras-yolo3 (https://github.com/qqwweee/keras-yolo3), use git to clone or directly download it as a zip compressed file.

(3) Import PyCharm
Open PyCharm, create a new project, and import the source code of keras-yolo3 into PyCharm

4. Download the pre-training model The
YOLO official website provides the weight file trained by the YOLOv3 model, download and save it to the computer. The download address is https://pjreddie.com/media/files/yolov3.weights

5. Training the model
Next comes the crucial step: training the model. Before training the model, there are several preparations to do.
(1) Convert
the annotation data file used by YOLO. Each line consists of the path of the file, the position of the annotation frame (upper left corner, lower right corner), and the category ID. The format is: image_file_path x_min, y_min, x_max, y_max, class_id
An example is as follows:

This file format is different from the previously prepared VOC_2007 annotation file format. Keras-yolo3 provides a conversion script voc_annotation.py
to convert voc format to yolo format. Before converting the format, first open the voc_annotation.py file and modify The value of classes inside. For example, the panda object marked in voc_2007 in this case is named panda, so voc_annotation.py is modified to:

import xml.etree.ElementTree as ET
from os import getcwd

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]

classes = ["panda"]

Create a new folder VOCdevkit/VOC2007, put the annotation data folders Annotations, ImageSets, JPEGImages of pandas into the folder VOCdevkit/VOC2007, and then execute the conversion script, the code is as follows:

mkdir VOCdevkit
mkdir VOCdevkit/VOC2007
mv Annotations VOCdevkit/VOC2007
mv ImageSets VOCdevkit/VOC2007
mv JPEGImages VOCdevkit/VOC2007

source activate tensorflow
python voc_annotation.py

After conversion, files in yolo format will be automatically generated, including training set, test set, and validation set.

(2) Create a category file
In the keras-yolo3 source code imported by PyCharm, create a new category file my_class.txt in the model_data directory, and write the category of the marked object in it, one category per line, as follows:

(3) Convert the weight file
Convert the yolo weight file yolov3.weights downloaded earlier into a model file suitable for Keras. The conversion code is as follows:

source activate tensorflow
python convert.py -w yolov3.cfg yolov3.weights model_data/yolo_weights.h5

(4) Modify the path configuration of the training file
Modify the relevant path configuration in train.py, mainly including: annotation_path, classes_path, weights_path

Among them, the default batch_size in train.py is 32 (line 57), which refers to the number of batches processed each time. The larger the value, the higher the performance requirements of the machine, so it can be adjusted according to the actual situation of the computer. Low

(5) Training model
After the above configuration, everything is finally ready, and the training can be started by executing train.py.

After training, the default save path is logs/000/trained_weights_final.h5, which can be modified as needed. It is located in line 85 of train.py, and the name of the model saved can be modified.

6. Use the model
After completing the training of the model, call yolo.py to use our trained model.
First, modify the model path and category file path in yolo.py as follows:

class YOLO(object):
    _defaults = {
        "model_path": 'logs/000/trained_weights_final.h5',
        "anchors_path": 'model_data/yolo_anchors.txt',
        "classes_path": 'model_data/my_classes.txt',
        "score" : 0.3,
        "iou" : 0.45,
        "model_image_size" : (416, 416),
        "gpu_num" : 1,
    }

You can use the YOLO model by calling the YOLO class. For the convenience of testing, add the following code at the end of yolo.py. After modifying the image path, you can use your own yolo model

if __name__ == '__main__':
    yolo=YOLO()
    path = '/data/work/tensorflow/data/panda_test/1.jpg'
    try:
        image = Image.open(path)
    except:
        print('Open Error! Try again!')
    else:
        r_image, _ = yolo.detect_image(image)
        r_image.show()

    yolo.close_session()

After the execution, the cute panda was obediently circled, huh, huh

Through the above steps, we have learned to train our own target detection model based on YOLO, so that we can use SSD and YOLO to train our own data in combination with actual needs in the application, and select a target detection model with better effect.

Follow my official account "Big Data and Artificial Intelligence Lab" (BigdataAILab), and then reply to the " code " keyword to get the complete source code .

Recommended related reading

[AI combat] Hands-on training of your own target detection model (YOLO)

Guess you like