一、数据集下载
二、数据集划分
在模型训练时,将训练数据集分成
训练集
和验证集
, 测试集作为测试数据使用,下面是训练集划分代码:
import os
import random
import shutil
# 定义原始训练集路径
original_train_dir = 'GTSRB/Final_Training/Images'
# 定义新的训练集和验证集路径
new_train_dir = 'GTSRB/train'
valid_dir = 'GTSRB/valid'
# 创建新的训练集和验证集文件夹
os.makedirs(new_train_dir, exist_ok=True)
os.makedirs(valid_dir, exist_ok=True)
# 定义划分比例,这里将 20% 作为验证集
validation_split = 0.2
# 遍历每个类别文件夹
for class_folder in os.listdir(original_train_dir):
class_path = os.path.join(original_train_dir, class_folder)
if os.path.isdir(class_path):
# 创建新的训练集和验证集类别文件夹
new_train_class_dir = os.path.join(new_train_dir, class_folder)
valid_class_dir = os.path.join(valid_dir, class_folder)
os.makedirs(new_train_class_dir, exist_ok=True)
os.makedirs(valid_class_dir, exist_ok=True)
# 获取该类别下的所有图像文件
image_files = [f for f in os.listdir(class_path) if f.endswith('.ppm')]
# 打乱图像文件列表
random.shuffle(image_files)
# 计算验证集的数量
num_validation = int(len(image_files) * validation_split)
# 划分训练集和验证集
validation_files = image_files[:num_validation]
training_files = image_files[num_validation:]
# 移动或复制文件到新的训练集和验证集文件夹
for file in training_files:
src = os.path.join(class_path, file)
dst = os.path.join(new_train_class_dir, file)
shutil.copy(src, dst)
for file in validation_files:
src = os.path.join(class_path, file)
dst = os.path.join(valid_class_dir, file)
shutil.copy(src, dst)
print("数据划分完成。")
END