代碼使用方法注意修改一下路徑、驗證集比例、類別名稱，其他均不需要改動，自動劃分訓練集、驗證集、建好全部文件夾、一鍵自動生成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")