一些小修改

contrast learning/split.py

@@ -0,0 +1,111 @@
+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)