一、数据集下载

• 官网地址：附含数据集说明文档
• 点击下载：训练数据集
• 点击下载：测试数据集

二、数据集划分

在模型训练时，将训练数据集分成训练集和验证集， 测试集作为测试数据使用，下面是训练集划分代码：

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