散称训练数据前置处理

data_preprocessing/data_split.py

@@ -0,0 +1,73 @@
+import os
+import shutil
+import random
+from pathlib import Path
+
+def is_image_file(filename):
+    """检查文件是否为图像文件"""
+    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff')
+    return filename.lower().endswith(image_extensions)
+
+def split_directory(src_dir, train_dir, val_dir, split_ratio=0.9):
+    """
+    分割目录中的图像文件到train和val目录
+    :param src_dir: 源目录路径
+    :param train_dir: 训练集目录路径
+    :param val_dir: 验证集目录路径
+    :param split_ratio: 训练集比例(默认0.9)
+    """
+    # 创建目标目录
+    Path(train_dir).mkdir(parents=True, exist_ok=True)
+    Path(val_dir).mkdir(parents=True, exist_ok=True)
+
+    # 遍历源目录
+    for root, dirs, files in os.walk(src_dir):
+        # 获取相对路径(相对于src_dir)
+        rel_path = os.path.relpath(root, src_dir)
+
+        # 跳过当前目录(.)
+        if rel_path == '.':
+            continue
+
+        # 创建对应的目标子目录
+        train_subdir = os.path.join(train_dir, rel_path)
+        val_subdir = os.path.join(val_dir, rel_path)
+        os.makedirs(train_subdir, exist_ok=True)
+        os.makedirs(val_subdir, exist_ok=True)
+
+        # 筛选图像文件
+        image_files = [f for f in files if is_image_file(f)]
+        if not image_files:
+            continue
+
+        # 随机打乱文件列表
+        random.shuffle(image_files)
+
+        # 计算分割点
+        split_point = int(len(image_files) * split_ratio)
+
+        # 复制文件到训练集
+        for file in image_files[:split_point]:
+            src = os.path.join(root, file)
+            dst = os.path.join(train_subdir, file)
+            # shutil.copy2(src, dst)
+            shutil.move(src, dst)
+
+        # 复制文件到验证集
+        for file in image_files[split_point:]:
+            src = os.path.join(root, file)
+            dst = os.path.join(val_subdir, file)
+            # shutil.copy2(src, dst)
+            shutil.move(src, dst)
+
+        print(f"处理完成: {rel_path} (共 {len(image_files)} 个图像, 训练集: {split_point}, 验证集: {len(image_files)-split_point})")
+
+if __name__ == "__main__":
+    # 设置目录路径
+    SOURCE_DIR = "scatter_add"
+    TRAIN_DIR = "scatter_data/train"
+    VAL_DIR = "scatter_data/val"
+
+    print("开始分割数据集...")
+    split_directory(SOURCE_DIR, TRAIN_DIR, VAL_DIR)
+    print("数据集分割完成")