如何用sklearn/train_test_split处理多级子类结构的CNN数据集划分？

针对你这种嵌套类别的数据集（水果类别下嵌套成熟度子类），直接用train_test_split肯定会失败——它只接受扁平化的样本列表和标签。解决思路是先把嵌套结构转换成扁平的标签体系，再进行划分，具体步骤和代码如下：

1. 收集所有样本的路径与标签 #

先遍历训练集的嵌套文件夹，把每个图像的路径和复合类别标签（比如apple_ripe、pineapple_unripe）整理成两个列表，这是后续划分的基础。

import os
from sklearn.model_selection import train_test_split
import shutil

# 替换成你的数据集根目录
root_dir = "你的数据集根目录路径"
train_dir = os.path.join(root_dir, "Train")

image_paths = []
labels = []

# 遍历水果类别文件夹
for fruit_name in os.listdir(train_dir):
    fruit_folder = os.path.join(train_dir, fruit_name)
    if not os.path.isdir(fruit_folder):
        continue
    # 遍历该水果下的成熟度子类
    for ripeness_state in os.listdir(fruit_folder):
        ripeness_folder = os.path.join(fruit_folder, ripeness_state)
        if not os.path.isdir(ripeness_folder):
            continue
        # 收集所有图像文件
        for img_file in os.listdir(ripeness_folder):
            if img_file.lower().endswith(('.png', '.jpg', '.jpeg')):
                img_path = os.path.join(ripeness_folder, img_file)
                image_paths.append(img_path)
                # 生成复合标签，用下划线连接水果和成熟度
                labels.append(f"{fruit_name.lower()}_{ripeness_state.lower()}")

2. 用train_test_split分层划分数据集 #

一定要用stratify参数，保证划分后的训练/验证集里每个复合类别的比例和原数据一致，避免类别失衡影响模型训练。

# 划分训练集和验证集，验证集占比可自行调整
train_paths, val_paths, train_labels, val_labels = train_test_split(
    image_paths,
    labels,
    test_size=0.2,
    random_state=42,
    stratify=labels  # 分层抽样是关键
)

3. 加载划分后的数据（两种可选方案） #

方案一：整理成标准文件夹结构（适配ImageDataGenerator） #

如果习惯用Keras的ImageDataGenerator加载数据，可以把划分后的文件复制到新的标准文件夹结构中（每个复合类别一个文件夹）：

# 新数据集根目录，自行指定
new_dataset_root = "新的数据集保存路径"
new_train_dir = os.path.join(new_dataset_root, "train")
new_val_dir = os.path.join(new_dataset_root, "val")
new_test_dir = os.path.join(new_dataset_root, "test")

# 复制原测试集到新目录
shutil.copytree(os.path.join(root_dir, "Test"), new_test_dir)

# 复制训练集文件
for path, label in zip(train_paths, train_labels):
    target_folder = os.path.join(new_train_dir, label)
    os.makedirs(target_folder, exist_ok=True)
    shutil.copy(path, target_folder)

# 复制验证集文件
for path, label in zip(val_paths, val_labels):
    target_folder = os.path.join(new_val_dir, label)
    os.makedirs(target_folder, exist_ok=True)
    shutil.copy(path, target_folder)

之后就可以用flow_from_directory正常加载：

from tensorflow.keras.preprocessing.image import ImageDataGenerator

train_datagen = ImageDataGenerator(rescale=1./255)
val_datagen = ImageDataGenerator(rescale=1./255)
test_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
    new_train_dir,
    target_size=(224, 224),  # 匹配你的CNN输入尺寸
    batch_size=32,
    class_mode='categorical'
)

val_generator = val_datagen.flow_from_directory(
    new_val_dir,
    target_size=(224, 224),
    batch_size=32,
    class_mode='categorical'
)

方案二：用tf.data直接加载（无需复制文件） #

如果不想占用额外存储空间，用tf.data.Dataset直接从路径列表加载图像更高效：

import tensorflow as tf
from sklearn.preprocessing import LabelEncoder

# 把字符串标签转换成整数编码
label_encoder = LabelEncoder()
train_labels_encoded = label_encoder.fit_transform(train_labels)
val_labels_encoded = label_encoder.transform(val_labels)

# 定义图像加载函数
def load_preprocess_image(img_path, label):
    img = tf.io.read_file(img_path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, (224, 224))  # 调整尺寸
    img = img / 255.0  # 归一化
    return img, label

# 创建训练集数据集
train_dataset = tf.data.Dataset.from_tensor_slices((train_paths, train_labels_encoded))
train_dataset = train_dataset.map(load_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)
train_dataset = train_dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)

# 创建验证集数据集
val_dataset = tf.data.Dataset.from_tensor_slices((val_paths, val_labels_encoded))
val_dataset = val_dataset.map(load_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)
val_dataset = val_dataset.batch(32).prefetch(tf.data.AUTOTUNE)

关键说明 #

• 之前划分失败的核心原因是train_test_split无法识别嵌套文件夹结构，必须先将多层类别合并为单一的复合标签，把数据扁平化。
• 分层抽样（stratify=labels）必须开启，否则小样本类别可能在划分后集中在某一个子集里，导致模型训练效果差。

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