YOLO类数据集划分代码
import os
import random
import shutil
# 定义文件路径
image_folder = "F:/Data_increase/images"
annotation_folder = "F:/Data_increase/yolo_txt"
train_folder = "F:/Data_increase/train"
train_xml_folder = "F:/Data_increase/train_txt"
val_folder = "F:/Data_increase/val"
val_xml_folder = "F:/Data_increase/val_txt"
test_folder = "F:/Data_increase/test"
test_xml_folder = "F:/Data_increase/test_txt"
# 创建保存文件夹
os.makedirs(train_folder, exist_ok=True)
os.makedirs(train_xml_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)
os.makedirs(val_xml_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)
os.makedirs(test_xml_folder, exist_ok=True)
# 获取所有图像文件的路径
image_files = [f for f in os.listdir(image_folder) if f.endswith('.jpg')]
# 随机打乱图像文件列表
random.shuffle(image_files)
# 计算划分数量
total_images = len(image_files)
train_split = int(0.7 * total_images)
val_split = int(0.2 * total_images)
# 划分并复制图像和对应的标注文件
for i, image_file in enumerate(image_files):
annotation_file = image_file.replace('.jpg', '.txt')
if i < train_split:
shutil.copy(os.path.join(image_folder, image_file), os.path.join(train_folder, image_file))
shutil.copy(os.path.join(annotation_folder, annotation_file), os.path.join(train_xml_folder, annotation_file))
elif i < train_split + val_split:
shutil.copy(os.path.join(image_folder, image_file), os.path.join(val_folder, image_file))
shutil.copy(os.path.join(annotation_folder, annotation_file), os.path.join(val_xml_folder, annotation_file))
else:
shutil.copy(os.path.join(image_folder, image_file), os.path.join(test_folder, image_file))
shutil.copy(os.path.join(annotation_folder, annotation_file), os.path.join(test_xml_folder, annotation_file))
|