112 lines
4.2 KiB
Python
112 lines
4.2 KiB
Python
import os
|
||
import shutil
|
||
import random
|
||
from pathlib import Path
|
||
|
||
|
||
def split_dataset(source_dir: str, output_dir: str, test_ratio: float = 0.2):
|
||
"""
|
||
将源文件夹中的文件按比例划分到输出文件夹下的 train 和 test 子目录中。
|
||
|
||
Args:
|
||
source_dir (str): 包含所有数据文件的源文件夹路径。
|
||
output_dir (str): 用于存放'train'和'test'文件夹的目标文件夹路径。
|
||
test_ratio (float, optional): 测试集所占的比例。默认为 0.2。
|
||
"""
|
||
print("--- 开始执行数据集划分 ---")
|
||
|
||
# 1. 路径处理和验证
|
||
source_path = Path(source_dir)
|
||
output_path = Path(output_dir)
|
||
|
||
if not source_path.is_dir():
|
||
print(f"错误：源文件夹 '{source_dir}' 不存在或不是一个目录。")
|
||
return
|
||
|
||
# 2. 创建输出文件夹 (train 和 test)
|
||
train_dir = output_path / 'train'
|
||
test_dir = output_path / 'test'
|
||
|
||
try:
|
||
os.makedirs(train_dir, exist_ok=True)
|
||
os.makedirs(test_dir, exist_ok=True)
|
||
print(f"输出目录已准备好: \n 训练集 -> {train_dir}\n 测试集 -> {test_dir}")
|
||
except OSError as e:
|
||
print(f"错误：创建输出目录时发生错误: {e}")
|
||
return
|
||
|
||
# 3. 获取所有文件并随机打乱
|
||
all_files = [f for f in source_path.iterdir() if f.is_file()]
|
||
|
||
if not all_files:
|
||
print(f"警告：源文件夹 '{source_dir}' 中没有文件可供划分。")
|
||
return
|
||
|
||
random.shuffle(all_files)
|
||
total_files = len(all_files)
|
||
print(f"在源文件夹中找到 {total_files} 个文件。")
|
||
|
||
# 4. 计算分割数量
|
||
num_test = int(total_files * test_ratio)
|
||
num_train = total_files - num_test
|
||
|
||
print(f"划分计划 -> 训练集: {num_train} 个文件 | 测试集: {num_test} 个文件")
|
||
|
||
# 5. 分割文件列表
|
||
test_files = all_files[:num_test]
|
||
train_files = all_files[num_test:]
|
||
|
||
# 6. 定义一个复制/移动文件的辅助函数
|
||
def copy_files(files_to_copy, destination_dir):
|
||
copied_count = 0
|
||
for file_path in files_to_copy:
|
||
try:
|
||
# 注意：这里使用的是复制(copy)，更安全。
|
||
# 如果你确认要移动(move)并且清空源文件夹，请将 shutil.copy 改为 shutil.move
|
||
shutil.copy(file_path, destination_dir)
|
||
copied_count += 1
|
||
except Exception as e:
|
||
print(f"处理文件 '{file_path.name}' 时出错: {e}")
|
||
return copied_count
|
||
|
||
# 7. 复制文件到对应的文件夹
|
||
print(f"\n正在复制文件到 'train' 文件夹...")
|
||
copied_train = copy_files(train_files, train_dir)
|
||
print(f"成功复制 {copied_train} 个文件到训练集。")
|
||
|
||
print(f"\n正在复制文件到 'test' 文件夹...")
|
||
copied_test = copy_files(test_files, test_dir)
|
||
print(f"成功复制 {copied_test} 个文件到测试集。")
|
||
|
||
print("\n--- 数据集划分完成！ ---")
|
||
|
||
|
||
# --- 如何使用这个函数 ---
|
||
if __name__ == '__main__':
|
||
# --- 请在这里配置你的文件夹路径 ---
|
||
|
||
# 你的原始数据集所在的文件夹
|
||
# 例如: 'C:/Users/YourUser/Desktop/my_dataset' (Windows)
|
||
# 或: '/home/user/project/raw/all_images' (Linux/macOS)
|
||
SOURCE_DATA_DIR = 'D:/download/2025-10/data_OS/input/S'
|
||
|
||
# 你希望将'train'和'test'文件夹创建在哪里
|
||
# 例如: 'C:/Users/YourUser/Desktop/split_output' (Windows)
|
||
# 或: '/home/user/project/raw/processed' (Linux/macOS)
|
||
# 如果使用 '.', 表示在当前脚本所在的目录下创建
|
||
OUTPUT_DIR = 'D:/download/2025-10/data_OS/train/S'
|
||
|
||
# --- 配置完成，下面是调用函数 ---
|
||
|
||
# 检查示例路径是否存在，如果不存在则创建并填充一些假文件用于演示
|
||
if not os.path.exists(SOURCE_DATA_DIR):
|
||
print(f"演示目录 '{SOURCE_DATA_DIR}' 不存在，正在创建并生成100个示例文件...")
|
||
os.makedirs(SOURCE_DATA_DIR)
|
||
for i in range(100):
|
||
with open(os.path.join(SOURCE_DATA_DIR, f'file_{i + 1:03d}.txt'), 'w') as f:
|
||
f.write(f'This is file {i + 1}.')
|
||
print("示例文件创建完毕。")
|
||
|
||
# 调用函数执行划分
|
||
split_dataset(SOURCE_DATA_DIR, OUTPUT_DIR, test_ratio=0.2)
|