如何从CNN图像分类模型的预测结果中获取类别标签/索引

• 第一步：获取类别映射关系
  你使用flow_from_directory初始化训练/验证数据生成器时，生成器自带class_indices属性，存储了类别名称到整数索引的对应关系，你需要先拿到该映射：

# 训练阶段你应该写过类似的生成器初始化代码，直接取class_indices即可
from tensorflow.keras.preprocessing.image import ImageDataGenerator
train_datagen = ImageDataGenerator(rescale=1./255)
train_generator = train_datagen.flow_from_directory(
    './train_dataset_path', # 你的训练集路径，按类别分设子文件夹
    target_size=(224,224),
    batch_size=32,
    class_mode='categorical'
)
# class_indices格式示例：{'猫':0, '狗':1, '鸟':2}
class_indices = train_generator.class_indices
# 转换为索引到类别名的映射，方便后续查询
idx_to_class = {v:k for k, v in class_indices.items()}

• 第二步：处理预测输出得到类别索引
  你给出的预测输出是4维张量，形状为(样本数, 特征图高度, 特征图宽度, 类别数)，最后一维为3，对应3分类任务，分两种场景处理：

  • 普通图像分类场景：正常分类头输出应为(样本数, 类别数)的2维张量，你当前的输出应该是分类层前的特征图输出，先压缩空间维度再取最大置信度索引即可：

  import numpy as np
  # preds为你给出的预测输出数组
  # 对高度、宽度维度做全局平均池化，压缩空间维度
  pooled_preds = np.mean(preds, axis=(1,2))
  # 取每个样本最大置信度对应的索引，即预测类别索引
  pred_class_ids = np.argmax(pooled_preds, axis=1)
  

  • 逐像素分类/语义分割场景：直接在最后一维取最大值索引，得到每个像素对应的类别索引：

  # 输出形状为(样本数, 特征图高度, 特征图宽度)，每个值对应位置的类别索引
  pred_class_ids = np.argmax(preds, axis=-1)
  

• 第三步：索引转换为类别名称
  拿到预测的类别索引后，用之前生成的idx_to_class映射查询即可得到具体类别名：

# 普通分类场景，获取第一个样本的预测类别名
sample_pred_label = idx_to_class[pred_class_ids[0]]

# 分割场景如果需要整图分类标签，可以统计所有像素出现最多的类别索引再查询

内容的提问来源于stack exchange，提问作者data.is.world