yolov5 trains its own data set - a small note

Some records in the process of my own study, if there are any inaccuracies, I implore everyone to criticize and correct!

Dataset format conversion

I downloaded the data set (MAR20) in VOC format. First, I need to convert the data set into a format that yolo can use.

1. merge xml files

A picture of the MAR20 dataset corresponds to two xml files, first merge the xml files (due to my limited level, and the information I consulted is one picture corresponding to one xml file, so first merge the two xml files, but the two xml files The file corresponds to two marking methods, whether it can be combined is doubtful, I hope someone can solve my dilemma ), the code is as follows:

from xml.etree.ElementTree import ElementTree, Element, parse
import xml.etree.ElementTree as ET
import os
import shutil
hole_path = 'D:/ProgramData/code/datasets/MAR20/Annotations/Horizontal Bounding Boxes'
arm_path = 'D:/ProgramData/code/datasets/MAR20/Annotations/Oriented Bounding Boxes'
out_path = 'D:/ProgramData/code/datasets/MAR20/Annotations/xmlcombine'
# 格式化
def __indent(elem, level=0):
    i = "\n" + level*"\t"
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "\t"
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            __indent(elem, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i

for hole_xml in os.listdir(hole_path):
    # 将同名xml合并
    if os.path.exists(os.path.join(arm_path,hole_xml)):
        print('Horizontal Bounding Boxes',hole_xml)
        tree_hole = parse(os.path.join(hole_path,hole_xml))
        root_hole = tree_hole.getroot()  # annotation
        new_hole = tree_hole
        tree_arm = parse(os.path.join(arm_path,hole_xml))
        root_arm = tree_arm.getroot()  # annotation
        object = (tree_arm.findall('object'))
        for i in range(len(object)):
            root_hole.append(object[i])
        __indent(root_hole)
        new_hole.write(os.path.join(out_path,hole_xml))
#不同名xml复制
    else:
     print('copying',hole_xml)
     shutil.copy(os.path.join(hole_path,hole_xml), out_path)
# 将不同名xml复制
    for arm_xml in os.listdir(arm_path):
        if not os.path.exists(os.path.join(out_path,arm_xml)):
            print('copying')
            shutil.copy(os.path.join(arm_path, arm_xml), out_path

2.xml conversion txt

Create a new labels folder under the root directory of the dataset to store the files converted from xml files to txt format. The specific operation is as follows:

Create a new txt_write.py file in the root directory of the pycharm project to divide the data set. The code is as follows:

# -*- coding: utf-8 -*-
"""
分训练集、验证集和测试集，按照 8：1：1 的比例来分，训练集8，验证集1，测试集1

"""
import os
import random
import argparse

parser = argparse.ArgumentParser()
# xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='D:/ProgramData/code/datasets/MAR20/Annotations/Horizontal Oriented Bounding Boxes', type=str, help='input xml label path')
# 数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='D:/ProgramData/code/datasets/MAR20/ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()

train_percent = 0.8  # 训练集所占比例
val_percent = 0.1  # 验证集所占比例
test_persent = 0.1  # 测试集所占比例

xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)

if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_xml)
list = list(range(num))

t_train = int(num * train_percent)
t_val = int(num * val_percent)

train = random.sample(list, t_train)
num1 = len(train)
for i in range(num1):
    list.remove(train[i])

val_test = [i for i in list if not i in train]
val = random.sample(val_test, t_val)
num2 = len(val)
for i in range(num2):
    list.remove(val[i])

file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')

for i in train:
    name = total_xml[i][:-4] + '\n'
    file_train.write(name)

for i in val:
    name = total_xml[i][:-4] + '\n'
    file_val.write(name)

for i in list:
    name = total_xml[i][:-4] + '\n'
    file_test.write(name)

file_train.close()
file_val.close()
file_test.close()

Convert the xml file to txt format to generate txt for the training set, verification set and test set. txt records the path of the picture

# *coding:utf-8 *
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import cv2

sets = ['train', 'test','val']   
#数据集类别
classes = ['A1','A2','A3','A4','A5','A6','A7','A8','A9','A10','A11','A12','A13','A14','A15','A16','A17','A18','A19','A20' ]

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('D:/yolov7-main/voc2007/Annotations/%s.xml' % (image_id))  # 修改路径（最好使用全路径）
    # img_file = cv2.imread('JILI_NB\images\%s.jpg' % (image_id))
    # out_file = open('JILI_NB\labels/%s.txt' % (image_id), 'w+')  # 修改路径（最好使用全路径）
    in_file = open('D:/ProgramData/code/datasets/MAR20/Annotations/Horizontal Oriented Bounding Boxes/%s.xml' % (image_id))  # 修改路径（最好使用全路径）
    img_file = cv2.imread('D:/ProgramData/code/datasets/MAR20/JPEGImages/%s.jpg' % (image_id))
    out_file = open('D:/ProgramData/code/datasets/MAR20/Labels/%s.txt' % (image_id), 'w+')  # 修改路径（最好使用全路径）
    tree = ET.parse(in_file)
    root = tree.getroot()
    # size = root.find('size')
    assert img_file is not None
    size = img_file.shape[0:-1]
    h = int(size[0])
    w = int(size[1])
    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))
        ZIP_ONE = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in ZIP_ONE]) + '\n')

wd = getcwd()
for image_set in sets:
    # if not os.path.exists('JILI_NB/labels'):
    #     os.makedirs('JILI_NB/labels/')
    if not os.path.exists('D:/ProgramData/code/datasets/MAR20/Labels'):
        os.makedirs('D:/ProgramData/code/datasets/MAR20/Labels')
    image_ids = open('D:/ProgramData/code/datasets/MAR20/ImageSets/Main/%s.txt' % (image_set)).read().split()  # 修改路径（最好使用全路径）
    list_file = open('D:/ProgramData/code/datasets/MAR20/Labels%s.txt' % (image_set), 'w+')  # 修改路径（最好使用全路径）
    # print(image_ids)
    for image_id in image_ids:
        try:
            print(image_id)
            list_file.write('D:/ProgramData/code/datasets/MAR20/JPEGImages/%s.jpg\n' % (image_id))  # 修改路径（最好使用全路径）
            convert_annotation(image_id)
        except:
            print('error img:', image_id)
    list_file.close()

3. Change folder structure

After the above operation is completed, there is no problem with the file format of the dataset. Next, the folder structure needs to be changed.

After performing the above operations, the folder structure is as follows, further changes are required

Change it to the picture below, and the code is as follows: (Note here, shutil.copy() calls the path in the txt file. Since the generated txt is the name of the picture, not the path, the following code will always report that it cannot be found when running The fault of the file, it took a day to find the problem. Added a piece of code for this problem, and it can run normally after writing the path into the txt file)

import shutil
import os
 
file_List = ["test", "val","train" ]
for file in file_List:
    if not os.path.exists('D:/ProgramData/code/yolov5-master/MAR20/images/%s' % file):
        os.makedirs('D:/ProgramData/code/yolov5-master/MAR20/images/%s' % file)
    if not os.path.exists('D:/ProgramData/code/yolov5-master/MAR20/labels/%s' % file):
        os.makedirs('D:/ProgramData/code/yolov5-master/MAR20/labels/%s' % file)
    print(os.path.exists('../tmp/ImageSets/Main/%s.txt' % file))
    f = open('../tmp/ImageSets/Main/%s.txt' % file, 'r')
    lines = f.readlines()
 '''
#将路径写进txt文件，如果生成的txt文件为图片路径则不需要
    f.seek(0)
    f.truncate()  # 先将原来文件进行清空
    for line_list in lines:  # 对于原来文件的内容每一行进行添加的操作
        line_list = line_list.replace("\n", "")
        line_new ="D:\ProgramData\code\yolov5-master\tmp\JPEGImages\" + line_list + ".jpg "+"\n"
        f.write(line_new)
    f.close()
'''
for line in lines:
        print(line)
        line = "/".join(line.split('/')[-5:]).strip()
        shutil.copy(line, "../MAR20/images/%s" % file)
        #shutil.copyfile(line, "../MAR20/images/%s" % file)
        #shutil.copymode(line, "../MAR20/images/%s" % file)
        line = line.replace('JPEGImages', 'Labels')
        line = line.replace('jpg', 'txt')
        shutil.copy(line, "../MAR20/labels/%s/" % file)

At this point, the data set conversion is all completed!

configuration file

4. Create a new datastes folder in the same level folder of the pycharm project

The datasets folder is used to store datasets, put the generated dataset folder in the datasets folder, and then start configuring the file

First, create a new airplane.yaml file in the data folder under the project directory

train: D:\ProgramData\code\datasets\MAR20\images\train
val: D:\ProgramData\code\datasets\MAR20\images\val
test: D:\ProgramData\code\datasets\MAR20\images\test

nc: 20 #有多少个类
#类别名称
names: ['A1','A2','A3','A4','A5','A6','A7','A8','A9','A10','A11','A12','A13','A14','A15','A16','A17','A18','A19','A20']

Next, modify the train.py file in the project root directory

#yolov5有smlx四个模型，选择哪个模型设置哪个模型路径（所有的路径都最好是全路径，不用ROOT代替）
parser.add_argument('--cfg', type=str, default='D:/ProgramData/code/yolov5-master/models/yolov5s.yaml', help='model.yaml path')
#调用自己设置的yaml文件
parser.add_argument('--data', type=str, default= 'D:/ProgramData/code/yolov5-master/data/airplane.yaml', help='dataset.yaml path')

After these are completed, you can start training.

Extension: epochs: how many rounds of training, default 100, can be changed as needed.

batch-size: The number of samples selected for each training, for example, if I have 3073 samples, batch-size=16, then the number of samples for each training is 193.