• Reprint please note the original sources, thank: blog.csdn.net/pentiumCM/a…

Deep learning – YoloV5 transfer learning, custom data set training

A list,

This paper introduces how to train the YOLO5 target detection model on our own data sets, including VOC and COCO data sets.

2. Custom data set training

During yoloV5 training, the code recognizes images and labels folders by default. Images are used as the input image folder and labels as the label folder. Therefore, no matter what data set format is adopted, as long as these two folders can be generated eventually, it is enough.

(I) VOC format data set

1. Prepare data sets

The data set directory is as follows:

└ ─ cell_phone: Annotations custom Datagram ├─ XML ├─images ├ V5 Use this folder as input for training images ├─ ├─ Main ├─ ├─ general exercises ├─ Labels yolov5 Use this folder as a marker for exercisesCopy the code

Instead of the default VOC data format, we created images and Labels folders in this step. Images directly copy images from JPEGImages into the labels folder. Labels as the folder to save the label files later. The directory is as follows:

2. Divide the data set

TXT, val. TXT, trainval. TXT, test. TXT. TXT, val. TXT, trainval. TXT and test. TXT files can be generated in ImageSets/Main by executing the following code

#! /usr/bin/env python # encoding: utf-8 ''' @Author : pentiumCM @Email : [email protected] @Software: PyCharm @File : make_txt.py @Time : 2020/11/25 16:51 @desc : Partitioning data sets -- generating: Trainval.txt, train.txt, val.txt, Test. TXT is a partition of training set + validation set, training set, validation set, and test set in the Main subfolder under ImageSets. Import OS import random def ds_partition(annotation_filepath, ds_divide_save__path): Train. TXT, val. TXT, trainval. TXT, test. TXT :param annotation_filepath: Annotations: The path of the Annotations file Annotations :param ds_divide_save__path """ if not os.path.exists(ds_divide_save__path): Os. mkdir(ds_divide_save__path) # train_percent: Percentage of training sets to (training sets + verification sets) Ratio of (training set + verification set) to total data set. The proportion of test sets is: 1-trainval_percent trainval_percent = 1 temp_xml = os.listdir(annotation_filepath) total_xml = [] for xml in temp_xml: if xml.endswith(".xml"): total_xml.append(xml) num = len(total_xml) list = range(num) tv = int(num * trainval_percent) tr = int(tv * train_percent) trainval = random.sample(list, tv) train = random.sample(trainval, tr) print("train and val size", tv) print("traub suze", tr) ftrainval = open(os.path.join(ds_divide_save__path, 'trainval.txt'), 'w') ftest = open(os.path.join(ds_divide_save__path, 'test.txt'), 'w') ftrain = open(os.path.join(ds_divide_save__path, 'train.txt'), 'w') fval = open(os.path.join(ds_divide_save__path, 'val.txt'), 'w') for i in list: name = total_xml[i][:-4] + '\n' if i in trainval: ftrainval.write(name) if i in train: ftrain.write(name) else: fval.write(name) else: ftest.write(name) ftrainval.close() ftrain.close() fval.close() ftest.close() if __name__ == '__main__': Annotations = 'annotation_filepath'; /cell_phone/' Annotations/' dataset_name + 'ImageSets/Main/' ds_partition(annotation_filepath, divide_save_path)Copy the code

3. Process the annotation result

VOC labels are stored in XML files. The rules for VOC corresponding to a label box are: Xmin ymin xmax Ymax YOLO Labels are stored in TXT files. The rules for YOLO corresponding to a label box are: class center_x center_y width height

Lables folder and test. TXT, train. TXT, val. TXT files will be generated after executing the code:

  • Labels store VOC parsing into YOLO label results
  • The other three TXT files are the absolute path of the pictures in the corresponding data set
#! /usr/bin/env python # encoding: utf-8 ''' @Author : pentiumCM @Email : [email protected] @Software: PyCharm @File : voc2yolo_label.py @Time : 2020/11/25 16:54 @desc : Category file format conversion - Convert dataset annotation format to Yolo_TXT format, Import xml.etree.ElementTree as ET import OS sets = ['train', 'test', ['cell phone'] def convert(size, box): """ Change VOC label to YOLO label, i.e. Xyxy -> XYWH: Param size: Param box: XyXY :return: "" "dw = 1 / size [0] dh = 1 / size [1] # center coordinates x = box (box [0] + [1])/y = 2.0 + box (box [2] [3]) / 2.0 # high w = box [1] -box [0] h = box[3] -box [2] # normalize x = x * dw w = w * dw y = y * dh h = h * dh return (x, y, w, h) def convert_annotation(annotation_filepath, label_dir, image_id): """ The result of VOC annotation is converted to the result of YOLO data set annotation: Param Annotation_filepath: VOC annotation folder path: Param Label_DIR: Param image_id: filename :return: """ in_file = open(annotation_filepath + '%s.xml' % (image_id)) out_file = open(label_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') if __name__ == '__main__': # Absolute path of current file abs_path = os.getcwd() + '/' # Custom dataset folder name dataset_name = 'cell_phone' # VOC dataset annotation file annotation_filepath = Annotations = dataset_name + '/ImageSets/Main/' # coco training image input folder image_dir = Label_dir = dataset_name + '/images/' # label_dir = dataset_name + '/labels/' if not os.path.exists(label_dir): os.makedirs(label_dir) for image_set in sets: txtfile = txtfile_dir + '%s.txt' % (image_set) image_ids = open(txtfile).read().strip().split() list_file = open(dataset_name + '/%s.txt' % (image_set), 'w') for image_id in image_ids: list_file.write((abs_path + image_dir + '%s.jpg\n' % (image_id)).replace('\\', '/')) convert_annotation(annotation_filepath, label_dir, image_id) list_file.close()Copy the code

The labels file contents in the Labels folder after parsing are shown in the figure below. Each line corresponds to a label box. The first value of the line corresponds to the category number, and the following four values correspond to the coordinates of the center point and the width and height of the label box (normalized to 0~1).

TXT, train. TXT and val. TXT files are shown as the following figure, which is the path of pictures in each dataset:

4. Build the dataset configuration file

We need to create a new YAML configuration file corresponding to our custom VOC data set, such as vocds.yaml, which basically says:

  • Train, val is the TXT file path generated in step 3
  • Nc: Number of categories (+1 is not required)
  • Names: indicates the name of a category

The contents of the YAML file are as follows:

# train2017 and val2017 data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] #train2017: .. /VOC/images/train2017/ # 16551 images #val2017: .. VOC/images/val2017 / train # 4952 images: F: / develop_code/python/yolo/yolov5 / v5.0 / yolov5 / data/cell_phone/train. TXT val: F: / develop_code/python/yolo/yolov5 / v5.0 / yolov5 / data/cell_phone/val. TXT # number of classes: the category of the data set number of nc: 1 # class names names: ['cell phone']Copy the code

5. Customize the model profile

To train either model, modify the corresponding model.yaml. For example, the pre-training model we use is Yolov5s, so we need to modify the corresponding configuration file. I directly copy the corresponding configuration file and modify the number of classes predicted by the model. The number of classes is the number of classes in our user-defined data set

6. Modify training parameters

Modify train.py training parameters:

–data: corresponds to ‘data/cell_phone.yaml’

— CFG: corresponding to the pre-trained model configuration file, ‘models/cell_phone.yaml’

–weights: pre-training model weights file, such as yolov5s.pt

–epochs, — Batch-size: It depends on the graphics card configuration of your PC. If the graphics card is too large, it will report insufficient memory

7. Start training

  • Python train. Py. Py > train_nub.log 2> &1&nohup python train.py > train_nub.log 2> &1&.

  • . Cache: A cache file will be generated during the training to interrupt the training process and continue the training next time. (Also if we want to rewrite the training, we need to clean this file to prevent interference from the results of the last training)

  • The results of model training are in the runs folder in the root directory

  • Screenshot of training process:

8. Performance evaluation

Model performance evaluation:

# model test: python test. Py -- data. / data/vocds yaml - weights runs/train/exp5 / weights/best pt - augmentCopy the code


(2) Coco data set

  1. Prepare the dataset folder coco, which has the same code as yolov5. The directory structure is as follows:

    ├─ ├─ exercises - ├─ exercises - ├─ exercises - Exercises - Exercises - Exercises - Exercises - Exercises - Exercises - Exercises - Exercises - Exercises - Exercises - Exercises │ ├─ ├─ un├ ─ un├ ─ un├ ─ un├ ─ un├ ─ un├ ─ un├ ─ un├ ─ un├ ─ un├ ─ un├ ─ un└Copy the code

  2. Yaml create a new file called coco_custom.yaml that contains the following contents:

    # train2017 and val2017 data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] train: .. /coco/images/train2017/ val: .. /coco/images/val2017/ # number of classes nc: 2 # class names names: ['person_legal', 'person_illegal']Copy the code

  3. Modify the configuration file of the pretraining model

  4. Modify the training parameters of train.py

  5. Start training

  6. Model performance evaluation

The above content is the same as the steps in the previous section and will not be repeated here.

Analysis of training results

Yolov5 will generate expi directory in run directory every time train is completed (I stands for the number of results generated, exp0 is generated after the first training, exp1 is generated after the second training… Visual training results are stored in the expi directory: PNG (confusion matrix), F1_curve. PNG (F1 curve), p_curve. PNG (P curve), R_curve. PNG (R curve), PR_curve. PNG (PR curve), result.png, results.tx T and training weights (last-pt and best.pt); Here is a brief explanation of the parameters:

  • Result. The PNG:

The resources

  • www.freesion.com/article/772…
  • Zhouchen.blog.csdn.net/article/det…