如何用Keras/PyTorch加载嵌套目录数据并实现两级CNN分类

一、正确加载两级标签的数据集 #

flow_from_directory默认只识别当前目录下的直接子文件夹作为类别，所以你之前的写法只会读取Real和Fake，没法获取子目录的类别标签。下面给两种可行方案：

方案1：自定义数据生成器（灵活可控） #

手动遍历目录，为每张图片生成「真假」和「类别」两个标签：

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator

def custom_data_generator(data_dir, img_size, batch_size, datagen=None):
    img_paths = []
    labels_real_fake = []
    labels_class = []
    
    # 标签映射
    class_map = {'c1': 0, 'c2': 1, 'c3': 2, 'c4': 3}
    real_fake_map = {'Real': 0, 'Fake': 1}
    
    # 遍历所有图片
    for real_fake_dir in os.listdir(data_dir):
        if real_fake_dir not in real_fake_map:
            continue
        real_fake_label = real_fake_map[real_fake_dir]
        class_dirs = os.listdir(os.path.join(data_dir, real_fake_dir))
        for class_dir in class_dirs:
            if class_dir not in class_map:
                continue
            class_label = class_map[class_dir]
            img_files = [f for f in os.listdir(os.path.join(data_dir, real_fake_dir, class_dir)) if f.endswith(('.jpg','.png'))]
            for img_file in img_files:
                img_path = os.path.join(data_dir, real_fake_dir, class_dir, img_file)
                img_paths.append(img_path)
                labels_real_fake.append(real_fake_label)
                labels_class.append(class_label)
    
    # 打乱数据
    indices = np.random.permutation(len(img_paths))
    img_paths = np.array(img_paths)[indices]
    labels_real_fake = np.array(labels_real_fake)[indices]
    labels_class = np.array(labels_class)[indices]
    
    # 批量生成数据
    while True:
        for start in range(0, len(img_paths), batch_size):
            end = min(start + batch_size, len(img_paths))
            batch_paths = img_paths[start:end]
            batch_imgs = []
            for path in batch_paths:
                img = tf.keras.preprocessing.image.load_img(path, target_size=img_size)
                img_array = tf.keras.preprocessing.image.img_to_array(img)
                if datagen:
                    img_array = datagen.random_transform(img_array)
                # 根据你的模型选择预处理方式，这里用ResNet的示例
                img_array = tf.keras.applications.resnet50.preprocess_input(img_array)
                batch_imgs.append(img_array)
            batch_imgs = np.array(batch_imgs)
            # 返回两个标签的字典
            yield batch_imgs, {
                'real_fake_output': labels_real_fake[start:end],
                'class_output': labels_class[start:end]
            }

# 使用示例
train_datagen = ImageDataGenerator(rotation_range=20, width_shift_range=0.2, height_shift_range=0.2, horizontal_flip=True)
train_generator = custom_data_generator(
    data_dir='MainData',
    img_size=(224, 224),
    batch_size=32,
    datagen=train_datagen
)

方案2：修改目录结构（简单直接） #

把「真假+类别」的组合作为单独类别，比如Real_c1、Fake_c2，再通过生成器拆分标签：

1. 重新整理目录结构：

MainData/
    Real_c1/
    Real_c2/
    Real_c3/
    Real_c4/
    Fake_c1/
    Fake_c2/
    Fake_c3/
    Fake_c4/

2. 加载并拆分标签：

train_generator = train_datagen.flow_from_directory(
    'MainData',
    target_size=img_size,
    batch_size=batch_size,
    class_mode='categorical'
)

def split_labels_generator(generator):
    for imgs, combined_labels in generator:
        real_fake_labels = np.zeros((len(combined_labels), 2))
        class_labels = np.zeros((len(combined_labels), 4))
        for i, label in enumerate(combined_labels):
            class_idx = np.argmax(label)
            # 前4类是Real的c1-c4，后4类是Fake的c1-c4
            real_fake_labels[i, 0 if class_idx <4 else 1] = 1
            class_labels[i, class_idx %4] = 1
        yield imgs, {
            'real_fake_output': real_fake_labels,
            'class_output': class_labels
        }

train_generator = split_labels_generator(train_generator)

二、构建两级分类的CNN模型 #

采用多输出模型，共享特征提取层，同时输出两个分类结果：

from tensorflow.keras import layers, Model

# 输入层
input_layer = layers.Input(shape=(224, 224, 3))

# 特征提取层（可以替换成预训练模型，比如ResNet50）
x = layers.Conv2D(32, (3,3), activation='relu')(input_layer)
x = layers.MaxPooling2D((2,2))(x)
x = layers.Conv2D(64, (3,3), activation='relu')(x)
x = layers.MaxPooling2D((2,2))(x)
x = layers.Conv2D(128, (3,3), activation='relu')(x)
x = layers.MaxPooling2D((2,2))(x)
x = layers.Flatten()(x)
x = layers.Dense(256, activation='relu')(x)

# 真假分类分支
real_fake_output = layers.Dense(2, activation='softmax', name='real_fake_output')(x)
# 类别分类分支
class_output = layers.Dense(4, activation='softmax', name='class_output')(x)

# 定义模型
model = Model(inputs=input_layer, outputs=[real_fake_output, class_output])

# 编译模型，指定双损失和评估指标
model.compile(
    optimizer='adam',
    loss={
        'real_fake_output': 'sparse_categorical_crossentropy', # 如果用独热标签就换成categorical_crossentropy
        'class_output': 'sparse_categorical_crossentropy'
    },
    metrics={
        'real_fake_output': 'accuracy',
        'class_output': 'accuracy'
    }
)

# 训练模型，注意steps_per_epoch要和生成器的批次数量匹配
model.fit(
    train_generator,
    epochs=20,
    steps_per_epoch=len(img_paths)//batch_size # 自定义生成器需要手动计算，或者用generator的length属性
)

三、预测阶段使用 #

预测时可以同时得到两个结果：

# 加载测试图片
img = tf.keras.preprocessing.image.load_img('test_image.jpg', target_size=(224,224))
img_array = tf.keras.preprocessing.image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0)
img_array = tf.keras.applications.resnet50.preprocess_input(img_array)

# 预测
real_fake_pred, class_pred = model.predict(img_array)
real_fake_result = 'Real' if np.argmax(real_fake_pred) == 0 else 'Fake'
class_result = ['c1','c2','c3','c4'][np.argmax(class_pred)]

print(f"预测结果：{real_fake_result} - {class_result}")

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