文章目录
- 代码描述
- 代码
代码描述
把一个文件夹里的图片按照比例(下面设置为0.8)切分成train和test数据,分别放在两个文件夹。
代码
import os
import random
import shutil
from tqdm import tqdm
# 图片文件夹路径
image_folder = './images'
# 训练文件夹路径
train_folder = 'train'
# 测试文件夹路径
test_folder = 'test'
# 确保训练和测试文件夹存在
os.makedirs(train_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)
# 列出所有图片文件
image_files = os.listdir(image_folder)
# 计算训练集和测试集的大小
train_size = int(len(image_files) * 0.8)
test_size = len(image_files) - train_size
# 随机打乱图片文件列表
random.shuffle(image_files)
# 将文件分成训练和测试集
train_files = image_files[:train_size]
test_files = image_files[train_size:]
# 复制训练集图片到训练文件夹
for filename in tqdm(train_files,total=len(train_files)):
src_path = os.path.join(image_folder, filename)
dst_path = os.path.join(train_folder, filename)
shutil.copy(src_path, dst_path)
# 复制测试集图片到测试文件夹
for filename in tqdm(test_files,total=len(test_files)):
src_path = os.path.join(image_folder, filename)
dst_path = os.path.join(test_folder, filename)
shutil.copy(src_path, dst_path)
print("图片已成功分割为训练集和测试集!")