写一个python脚本,将一个路径下的277个类别的文件夹氛围训练集(train)和验证集(val),两者数量之比为8:2,即最后该路径下应该只有2个文件夹,一个是train,里面有277个分类的图片(总数为80%),另一个是val,里面是277个分类的图片(总数为20%)
时间: 2024-11-07 15:12:57 浏览: 12
python脚本 自动切分data为val集 train集 test集
你可以使用Python的os模块来遍历目录结构,并利用shutil库来复制文件。以下是一个简单的脚本示例,它假设每个类别都有一个单独的文件夹,并且所有的图片都是jpg格式:
```python
import os
import shutil
# 定义源目录、目标训练集和验证集目录
src_dir = 'your_source_directory' # 替换为实际源目录
train_dir = 'train'
val_dir = 'val'
# 创建目标目录
if not os.path.exists(train_dir):
os.makedirs(train_dir)
if not os.path.exists(val_dir):
os.makedirs(val_dir)
# 获取所有类别名
category_folders = [d for d in os.listdir(src_dir) if os.path.isdir(os.path.join(src_dir, d))]
# 计算训练集和验证集的文件数
total_files = len(category_folders)
train_files = total_files * 0.8 # 训练集占总文件数的比例
val_files = total_files - train_files # 验证集剩余文件数
# 按比例分配文件到训练集和验证集
for category_idx, category_folder in enumerate(sorted(category_folders)):
src_category_dir = os.path.join(src_dir, category_folder)
files_in_category = [f for f in os.listdir(src_category_dir) if f.endswith('.jpg')] # 假设所有图片都是jpg格式
# 确保图片数量足够
assert len(files_in_category) >= val_files, "Not enough images in category {}, consider using a larger validation set ratio.".format(category_folder)
# 分配文件
num_train_files = int(len(files_in_category) * 0.8)
num_val_files = len(files_in_category) - num_train_files
train_path = os.path.join(train_dir, category_folder)
val_path = os.path.join(val_dir, category_folder)
if not os.path.exists(train_path):
os.makedirs(train_path)
if not os.path.exists(val_path):
os.makedirs(val_path)
# 将文件移动到对应的文件夹
shutil.move(os.path.join(src_category_dir, files_in_category[:num_train_files]), train_path)
shutil.move(os.path.join(src_category_dir, files_in_category[num_train_files:]), val_path)
阅读全文