如何理解image_dataset_from_directory()并提取X、y用于ROC曲线、混淆矩阵计算

你通过image_dataset_from_directory生成的是tf.data.Dataset类型的批次式数据集，要获取用于绘制ROC曲线、混淆矩阵的输入数据，可按照以下步骤操作：

1. 提取数据集的真实标签 #

不需要提前把所有图像读入内存，直接遍历数据集收集真实标签即可，两种实现方式：

import numpy as np

# 以测试集test_ds为例，验证集val_ds操作完全一致
# 方式1：遍历批次收集标签
y_true = []
for images, labels in test_ds:
    y_true.extend(labels.numpy())
y_true = np.array(y_true)

# 方式2：解批后直接提取，代码更简洁
y_true = np.array([label.numpy() for _, label in test_ds.unbatch()])

如果确实需要提取所有图像特征（仅建议小数据集使用，大数据集全量加载会占满内存），可使用如下代码：

X = np.array([img.numpy() for img, _ in test_ds.unbatch()])

2. 获取模型预测结果 #

直接调用模型的predict方法传入数据集即可得到所有样本的预测结果，不需要单独提取X：

# 得到所有样本的预测概率
y_pred_proba = model.predict(test_ds)

# 二分类场景：提取正类概率用于ROC计算，取预测类别用于混淆矩阵计算
y_pred_proba_pos = y_pred_proba[:, 1] if y_pred_proba.shape[1] > 1 else y_pred_proba
y_pred = np.argmax(y_pred_proba, axis=1) if y_pred_proba.shape[1] > 1 else np.round(y_pred_proba).astype(int)

3. 传入评估函数绘图 #

拿到真实标签y_true、预测类别y_pred、预测概率y_pred_proba后，直接传入对应的评估函数即可：

from sklearn.metrics import confusion_matrix, roc_curve, auc

# 计算混淆矩阵
cm = confusion_matrix(y_true, y_pred)
# 二分类场景计算ROC和AUC
fpr, tpr, thresholds = roc_curve(y_true, y_pred_proba_pos)
roc_auc = auc(fpr, tpr)

注意：即使你的数据集开启了shuffle也不会出现标签和预测结果不匹配的问题，predict遍历数据集的顺序和手动遍历数据集取标签的顺序完全一致。

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