如何从ImageDataGenerator.flow_from_directory提取单张图像与标签

嘿，我明白你的困惑了！你看到的(32,38)形状的数组是因为flow_from_directory默认的几个参数导致的：

• 默认batch_size=32，所以每次迭代返回一个包含32张图像的批次；
• 默认class_mode='categorical'，标签会被处理成one-hot编码的形式（38列对应38个类别，只有对应类别的位置是1，其余是0）。

下面给你两种简单的解决方案，按需选择：

方法1：修改class_mode为'int'，直接获取类别索引 #

这种方法最直观，标签会以整数形式返回，直接对应类别的索引，还可以结合class_indices拿到类别名称：

# 初始化生成器时修改class_mode，可选batch_size=1每次取一张
valid = ImageDataGenerator().flow_from_directory(
    directory="dataset/valid",
    target_size=(224,224),
    class_mode='int',  # 关键修改：标签为整数索引
    batch_size=1       # 可选：每次只返回1张图和1个标签
)

# 提取单张图像和标签
for img, lbl in valid:
    # 因为batch_size=1，所以lbl是形状(1,)的数组，取第一个元素就是索引
    class_idx = lbl[0]
    print(f"单张图像的类别索引: {class_idx}")
    
    # 如果需要类别名称，用class_indices的逆映射
    class_name_map = {v: k for k, v in valid.class_indices.items()}
    print(f"对应的类别名称: {class_name_map[class_idx]}")
    break

方法2：保留默认设置，从one-hot编码中提取索引 #

如果不想修改生成器的参数，可以用np.argmax()从one-hot数组中找到类别对应的索引：

import numpy as np

valid = ImageDataGenerator().flow_from_directory(
    directory="dataset/valid",
    target_size=(224,224)
)

# 遍历批次中的每张图像
for imgs_batch, lbls_batch in valid:
    # 取批次中的第一张图像和对应的one-hot标签
    single_img = imgs_batch[0]
    single_lbl_onehot = lbls_batch[0]
    
    # 从one-hot数组中获取类别索引
    class_idx = np.argmax(single_lbl_onehot)
    print(f"单张图像的类别索引: {class_idx}")
    
    # 同样可以映射到类别名称
    class_name_map = {v: k for k, v in valid.class_indices.items()}
    print(f"对应的类别名称: {class_name_map[class_idx]}")
    break

小提示 #

• valid.class_indices是生成器自带的字典，键是你的数据集文件夹名称（类别名），值是对应的整数索引，非常方便用来做索引和类别名的转换；
• 如果需要逐个处理所有验证集图像，把batch_size设为1，然后去掉break循环整个生成器即可。

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