74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
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("数据集分割完成")
|