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

目录

使用方法：

 全部代码：

---

使用方法：

 全部代码：

import os,random,shutilimport xml.etree.ElementTree as ET
import pickle
import os
from os import listdir,getcwd
from os.path import joindef convert(size,box):x_center=(box[0]+box[1])/2.0y_center=(box[2]+box[3])/2.0x=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').textcls=obj.find('name').textif cls not in classes or int(difficult)==1:continuecls_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)returndef 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改成.xmlif 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 folderos.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 folderos.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 folderos.makedirs("./val")if not os.path.exists("./train"):# Create the folderos.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 folderos.makedirs("./images")if not os.path.exists("./labels"):# Create the folderos.makedirs("./labels")#将所有文件移动到最终的文件夹中import shutil# Define the source and destination folderssource_folder = "./train"destination_folder = "./labels"# Move the files from the source folder to the destination foldershutil.move(source_folder, destination_folder)source_folder = "./val"destination_folder = "./labels"# Move the files from the source folder to the destination foldershutil.move(source_folder, destination_folder)source_folder = "./JPEGImages"destination_folder = "./images"# Move the files from the source folder to the destination foldershutil.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 foldershutil.move(source_folder, destination_folder)os.rename("./images/JPEGImages_val", "./images/val")