代码使用方法注意修改一下路径、验证集比例、类别名称，其他均不需要改动，自动划分训练集、验证集、建好全部文件夹、一键自动生成Yolo格式数据集在当前目录下，大家可以直接修改相应的配置文件进行训练。

目录

使用方法：

 全部代码：

---

使用方法：

 全部代码：

import os,random,shutil

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir,getcwd
from os.path import join


def convert(size,box):
    x_center=(box[0]+box[1])/2.0
    y_center=(box[2]+box[3])/2.0

    x=x_center/size[0]
    y=y_center/size[1]
    w=(box[1]-box[0])/size[0]
    h=(box[3]-box[2])/size[1]
    return (x,y,w,h)

def convert_annotation(xml_file_path,save_txt_file_path,classes):
    xml_file=os.listdir(xml_file_path)
    print(xml_file)
    for xml_name in xml_file:
        print(xml_file)
        xml_file=os.path.join(xml_file_path,xml_name)
        out_txt_path=os.path.join(save_txt_file_path,xml_name.split('.')[0]+".txt")
        out_txt_f=open(out_txt_path,'w')
        tree=ET.parse(xml_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))

            print(w,h,b)
            bb=convert((w,h),b)
            out_txt_f.write(str(cls_id)+" "+" ".join([str(a) for a in bb])+"\n")


def moving(fileDir,tarDir,rate=0.2):
    pathDir=os.listdir(fileDir)
    filenumber=len(pathDir)
                     #自定义抽取图片比例
    picknumber=int(filenumber*rate)#按照rate比例从文件夹中取一定数量的图片
    sample=random.sample(pathDir,picknumber)#随机选取picknumber数量的图片
    print(sample)

    for name in sample:
        shutil.move(fileDir+"/"+name,tarDir+"/"+name)
    return

def movelabel(file_list,file_label_train,file_label_val):
    for i in file_list:
        if i.endswith(".png") or i.endswith(".jpg"):
            #filename=file_label_train+"/"+i[:-4]       可以将.xml文件将.txt改成.xml文件
            filename=file_label_train+"/"+i[:-4]+".xml"       #可以改成xml文件将.txt改成.xml
            if os.path.exists(filename):
                shutil.move(filename,file_label_val)
                print("处理成功")

if __name__=="__main__":
    """
    设置图片路径、label路径、验证集比例、类别
    """
    fileDir=r"./JPEGImages"    #图片的路径
    file_label_train = r"./Annotations"    #标签文件的路径
    rate=0.2  #验证集的比例
    classes1 = ['fire']

    """
    以下均不需要改动
    """



    if not os.path.exists("./JPEGImages_val"):
        # Create the folder
        os.makedirs("./JPEGImages_val")

    tarDir=r"./JPEGImages_val"
    moving(fileDir,tarDir,rate)
    file_list=os.listdir(tarDir)

    if not os.path.exists("./Annotations_val"):
        # Create the folder
        os.makedirs("./Annotations_val")

    file_label_val=r"./Annotations_val"
    movelabel(file_list,file_label_train,file_label_val)

    #VOC转Yolo格式

    # 2 voc 格式的xml标签文件
    if not os.path.exists("./val"):
        # Create the folder
        os.makedirs("./val")

    if not os.path.exists("./train"):
        # Create the folder
        os.makedirs("./train")


    xml_files1 = r"./Annotations_val"
    save_txt_files1 = r"./val"

    convert_annotation(xml_files1, save_txt_files1, classes1)

    xml_files1 = r"./Annotations"
    save_txt_files1 = r"./train"

    convert_annotation(xml_files1, save_txt_files1, classes1)


    #创建所有文件夹
    if not os.path.exists("./images"):
        # Create the folder
        os.makedirs("./images")

    if not os.path.exists("./labels"):
        # Create the folder
        os.makedirs("./labels")


    #将所有文件移动到最终的文件夹中

    import shutil

    # Define the source and destination folders
    source_folder = "./train"
    destination_folder = "./labels"

    # Move the files from the source folder to the destination folder
    shutil.move(source_folder, destination_folder)


    source_folder = "./val"
    destination_folder = "./labels"

    # Move the files from the source folder to the destination folder
    shutil.move(source_folder, destination_folder)

    source_folder = "./JPEGImages"
    destination_folder = "./images"

    # Move the files from the source folder to the destination folder
    shutil.move(source_folder, destination_folder)
    os.rename("./images/JPEGImages", "./images/train")

    source_folder = "./JPEGImages_val"
    destination_folder = "./images"

    # Move the files from the source folder to the destination folder
    shutil.move(source_folder, destination_folder)
    os.rename("./images/JPEGImages_val", "./images/val")