将标签和对应的图像进行划分

import os
import random
import shutil

# 设置文件夹路径
images_folder = r'C:\Users\DELL\Desktop\mouse\output\images'  # 替换为实际的 images 文件夹路径
annotations_folder = r'C:\Users\DELL\Desktop\mouse\output\txt'  # 替换为实际的 txt_annotations 文件夹路径
output_folder = r'C:\Users\DELL\Desktop\mouse\output\split_data'  # 替换为保存划分结果的文件夹路径

# 设置划分比例
train_ratio = 0.8

# 确保输出文件夹存在
os.makedirs(output_folder, exist_ok=True)
train_images_folder = os.path.join(output_folder, 'train_images')
val_images_folder = os.path.join(output_folder, 'val_images')
train_annotations_folder = os.path.join(output_folder, 'train_annotations')
val_annotations_folder = os.path.join(output_folder, 'val_annotations')
os.makedirs(train_images_folder, exist_ok=True)
os.makedirs(val_images_folder, exist_ok=True)
os.makedirs(train_annotations_folder, exist_ok=True)
os.makedirs(val_annotations_folder, exist_ok=True)

# 获取所有图像文件
image_files = [f for f in os.listdir(images_folder) if f.endswith('.jpg') or f.endswith('.png')]
random.shuffle(image_files)

# 划分数据集
train_count = int(len(image_files) * train_ratio)
train_files = image_files[:train_count]
val_files = image_files[train_count:]

# 复制文件到对应的文件夹
for file in train_files:
    # 复制图像文件
    src_image_path = os.path.join(images_folder, file)
    dst_image_path = os.path.join(train_images_folder, file)
    shutil.copy(src_image_path, dst_image_path)

    # 复制对应的标签文件
    annotation_file = os.path.splitext(file)[0] + '.txt'
    src_annotation_path = os.path.join(annotations_folder, annotation_file)
    dst_annotation_path = os.path.join(train_annotations_folder, annotation_file)
    if os.path.exists(src_annotation_path):
        shutil.copy(src_annotation_path, dst_annotation_path)

for file in val_files:
    # 复制图像文件
    src_image_path = os.path.join(images_folder, file)
    dst_image_path = os.path.join(val_images_folder, file)
    shutil.copy(src_image_path, dst_image_path)

    # 复制对应的标签文件
    annotation_file = os.path.splitext(file)[0] + '.txt'
    src_annotation_path = os.path.join(annotations_folder, annotation_file)
    dst_annotation_path = os.path.join(val_annotations_folder, annotation_file)
    if os.path.exists(src_annotation_path):
        shutil.copy(src_annotation_path, dst_annotation_path)

print("数据集划分完成！")