代码使用方法注意修改一下路径、验证集比例、类别名称，其他均不需要改动，自动划分训练集、验证集、建好全部文件夹、一键自动生成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")
