yolov5训练自己的数据集(代码片段)

import os
import random
import shutil

img_path = \'img\'
xml_path = \'xml\'



def split_file_name(file_name):
    f_name, _ = file_name.split(\'.\')
    return f_name


def split_move_file(target_path, save_basic_path, train_scale=0.9):
    train_img_path = os.path.join(save_basic_path, \'images/train\')
    train_xml_path = os.path.join(save_basic_path, \'xml/train\')
    val_img_path = os.path.join(save_basic_path, \'images/val\')
    val_xml_path = os.path.join(save_basic_path, \'xml/val\')
    print(save_basic_path, train_img_path)
    if not os.path.exists(train_img_path):
        os.makedirs(train_img_path)
    if not os.path.exists(train_xml_path):
        os.makedirs(train_xml_path)
    if not os.path.exists(val_img_path):
        os.makedirs(val_img_path)
    if not os.path.exists(val_xml_path):
        os.makedirs(val_xml_path)


    img_file_path = os.path.join(target_path, img_path)
    file_list = os.listdir(img_file_path)
    # print(file_list)
    # 得到名字列表
    file_name_li = list(map(lambda x: split_file_name(x), file_list))
    random.shuffle(file_name_li)
    # print(file_name_li)
    train_ind = int(len(file_name_li) * train_scale)
    train_data = file_name_li[:train_ind]
    val_data = file_name_li[train_ind:]
    print(\'total number\', len(file_name_li))
    print(\'train number\', len(train_data))
    print(\'val number\', len(val_data))

    for file in train_data:

        file_path = os.path.join(img_file_path, file+\'.jpg\')
        save_path = os.path.join(train_img_path, file+\'.jpg\')
        if not os.path.exists(file_path):
            file_path = os.path.join(img_file_path, file + \'.jpeg\')
            save_path = os.path.join(train_img_path, file + \'.jpg\')
            if not os.path.exists(file_path):
                file_path = os.path.join(img_file_path, file + \'.png\')
                save_path = os.path.join(train_img_path, file + \'.png\')
        if os.path.exists(file_path):
            shutil.copyfile(file_path, save_path)

        # xml文件
        xml_file_path = os.path.join(target_path, xml_path)
        file_path = os.path.join(xml_file_path, file + \'.xml\')
        save_path = os.path.join(train_xml_path, file + \'.xml\')
        if os.path.exists(file_path):
            shutil.copyfile(file_path, save_path)
    for file in val_data:

        file_path = os.path.join(img_file_path, file+\'.jpg\')
        save_path = os.path.join(val_img_path, file+\'.jpg\')
        if not os.path.exists(file_path):
            file_path = os.path.join(img_file_path, file + \'.jpeg\')
            save_path = os.path.join(val_img_path, file + \'.jpg\')
            if not os.path.exists(file_path):
                file_path = os.path.join(img_file_path, file + \'.png\')
                save_path = os.path.join(val_img_path, file + \'.png\')
        if os.path.exists(file_path):
            shutil.copyfile(file_path, save_path)

        # xml文件
        xml_file_path = os.path.join(target_path, xml_path)
        file_path = os.path.join(xml_file_path, file + \'.xml\')
        save_path = os.path.join(val_xml_path, file + \'.xml\')
        if os.path.exists(file_path):
            shutil.copyfile(file_path, save_path)


if __name__ == \'__main__\':

    target_path = r\'C:\\Users\\mojia\\Desktop\\maizi\\maozi20230326\'
    save_basic_path = r\'C:\\Users\\mojia\\Desktop\\maizi\\maozi20230326_train_val\'
    if not os.path.exists(save_basic_path):
        os.mkdir(save_basic_path)
    scale = 0.9     # 训练集比例
    split_move_file(target_path, save_basic_path, scale)

import xml.etree.ElementTree as ET

import os





def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    if w >= 1:
        w = 0.99
    if h >= 1:
        h = 0.99
    return (x, y, w, h)

folder_li = [\'train\', \'val\']
def convert_annotation(rootpath, classes):
    labelpath = rootpath + \'/labels\'  # 生成的.txt文件会被保存在labels目录下
    if not os.path.exists(labelpath):
        os.makedirs(labelpath)
    for folder in folder_li:
        xmlpath = rootpath + \'/xml/\'+folder
        file_list = os.listdir(xmlpath)
        for xmlname in file_list:
            xmlfile = os.path.join(xmlpath, xmlname)
            with open(xmlfile, "r", encoding=\'UTF-8\') as in_file:
                txtname = xmlname[:-4] + \'.txt\'
                # print(txtname)
                txtpath = labelpath + \'/\' + folder
                if not os.path.exists(txtpath):
                    os.makedirs(txtpath)
                txtfile = os.path.join(txtpath, txtname)
                with open(txtfile, "w+", encoding=\'UTF-8\') as out_file:
                    tree = ET.parse(in_file)
                    root = tree.getroot()
                    size = root.find(\'size\')
                    w = int(size.find(\'width\').text)
                    h = int(size.find(\'height\').text)
                    out_file.truncate()
                    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__":
    rootpath = r\'C:\\Users\\mojia\\Desktop\\maizi\\maozi20230326_train_val\'
    # 数据标签
    classes = [\'帽子\']  # 需要修改
    convert_annotation(rootpath, classes)