1.需求
手动划分文件夹中的图片数据集为训练集、验证集和测试集,即进行文件夹中的数据集(都是图片)进行划分。
2.步骤
- 使用文件处理库(如os)遍历读取文件夹中的图片文件。
- 将读取到的图片文件路径存储到列表中。
- 打乱图片文件列表。
- 根据指定的比例或数量计算划分的样本数量。
- 利用切片操作将图片文件列表分为训练集、验证集和测试集。
- 将划分后的图片文件移动到相应的文件夹中。
3.代码实现
import os
import random
import shutil
def split_dataset_folder(dataset_folder, output_folder, train_ratio, val_ratio, test_ratio):
# 遍历读取文件夹中的图片文件
image_files = []
for root, dirs, files in os.walk(dataset_folder):
for file in files:
if file.endswith(".jpg") or file.endswith(".png"):
image_files.append(os.path.join(root, file))
# 打乱图片文件列表
random.shuffle(image_files)
# 计算划分的样本数量
dataset_size = len(image_files)
train_size = int(train_ratio * dataset_size)
val_size = int(val_ratio * dataset_size)
test_size = int(test_ratio * dataset_size)
# 划分图片文件列表
train_files = image_files[:train_size]
val_files = image_files[train_size:train_size+val_size]
test_files = image_files[train_size+val_size:train_size+val_size+test_size]
# 创建输出文件夹
os.makedirs(output_folder, exist_ok=True)
# 移动训练集图片文件
train_output_folder = os.path.join(output_folder, "train")
os.makedirs(train_output_folder, exist_ok=True)
for file in train_files:
shutil.move(file, train_output_folder)
# 移动验证集图片文件
val_output_folder = os.path.join(output_folder, "validation")
os.makedirs(val_output_folder, exist_ok=True)
for file in val_files:
shutil.move(file, val_output_folder)
# 移动测试集图片文件
test_output_folder = os.path.join(output_folder, "test")
os.makedirs(test_output_folder, exist_ok=True)
for file in test_files:
shutil.move(file, test_output_folder)
# 示例用法
dataset_folder = "path/to/your/dataset_folder" # 数据集文件夹路径
output_folder = "path/to/save/split_dataset" # 划分后的数据集存储路径
train_ratio = 0.7 # 训练集所占比例
val_ratio = 0.2 # 验证集所占比例
test_ratio = 0.1 # 测试集所占比例
split_dataset_folder(dataset_folder, output_folder, train_ratio, val_ratio, test_ratio)
4.注意事项
ataset_folder表示原始图片数据集所在的文件夹路径,output_folder表示划分后的数据集存储的文件夹路径,train_ratio、val_ratio和test_ratio表示训练集、验证集和测试集所占比例。代码将会遍历读取文件夹中的图片文件,打乱后按照比例划分为训练集、验证集和测试集,并将划分后的图片文件移动到相应的文件夹中。
实现效果
经过查看,文件夹里面的图像符合比例数量要求。