如何将按灾害类型分类的图像文件夹数据集处理后用于模型训练

1. 确认数据集结构 #

首先核对数据集的目录结构符合如下分类规则：

数据集根目录/
├── 火灾/
│ ├── img001.jpg
│ ├── img002.jpg
│ └── ...
├── 洪水/
│ ├── img101.jpg
│ └── ...
├── 地震/
└── 其他灾害分类文件夹...

根目录下的每个子文件夹名称对应一个灾害分类标签，子文件夹内存放该分类的所有灾害图像。

2. 生成CSV标注文件（非必需，按需选择） #

如果需要CSV格式的标注文件，可以通过遍历目录批量生成，Python示例代码如下：

import os
import csv

# 替换为你的数据集根目录路径
dataset_root = "./disaster_dataset"
# 替换为你要输出的CSV文件路径
output_label_path = "./dataset_labels.csv"
# 支持的图像后缀，可按需扩展
valid_img_suffix = ('.png', '.jpg', '.jpeg', '.bmp')

label_rows = []
# 遍历所有分类文件夹
for label_name in os.listdir(dataset_root):
    label_dir = os.path.join(dataset_root, label_name)
    if not os.path.isdir(label_dir):
        continue
    # 遍历分类下的所有图像
    for img_file in os.listdir(label_dir):
        if img_file.lower().endswith(valid_img_suffix):
            img_full_path = os.path.join(label_dir, img_file)
            label_rows.append([img_full_path, label_name])

# 写入CSV文件
with open(output_label_path, 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    # 表头可按需修改，也可以新增数字标签列，把字符串标签映射为0、1、2...的整数
    writer.writerow(['image_path', 'str_label'])
    writer.writerows(label_rows)

3. 数据集拆分 #

按照7:2:1或者8:1:1的比例将数据集拆分为训练集、验证集、测试集，拆分时需要保证每个分类的样本在三个集合中的占比和原数据集一致，避免标签分布偏移。
拆分后可以生成对应train.csv/val.csv/test.csv三个标注文件，也可以将拆分后的图像放到结构和原数据集一致的train/val/test三个根目录下。

4. 适配训练框架加载 #

主流深度学习框架都支持直接读取文件夹分类结构的数据集，无需提前生成CSV：

• PyTorch生态可以直接用torchvision.datasets.ImageFolder接口，传入数据集根目录后会自动识别文件夹名作为标签，可直接叠加数据增强的transform配置。
• TensorFlow/Keras生态可以直接用tf.keras.utils.image_dataset_from_directory接口，传入根目录即可自动完成数据加载和标签映射。

5. 加载后校验 #

完成加载后必须做以下校验：

• 核对加载的总样本数和原数据集的统计样本数一致
• 随机抽样部分样本，确认图像和对应标签匹配正确
• 提前过滤损坏、格式不兼容的图像，避免训练过程报错

内容的提问来源于stack exchange，提问作者anfal