TensorFlow加载数据集后获取class_indices报AttributeError的问题求助

问题场景 #

使用tf.keras.preprocessing.image_dataset_from_directory加载数据集后，尝试通过train_ds.class_indices获取类别与索引的映射时，出现报错：

AttributeError: '_BatchDataset' object has no attribute 'class_indices'

用户的数据集加载代码：

img_height= 220
img_width= 220
batch_size= 64
train_ds= tf.keras.preprocessing.image_dataset_from_directory(
    train_path,
    validation_split= 0.2,
    subset='training',
    seed=123,
    image_size=(img_height,img_width),
    batch_size= batch_size
)

val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    train_path,
    validation_split= 0.2,
    subset='validation',
    seed=123,
    image_size=(img_height,img_width),
    batch_size= batch_size,
    shuffle=False
)

test_ds = tf.keras.preprocessing.image_dataset_from_directory(
    test_path,
    seed=123,
    image_size=(img_height,img_width),
    batch_size= batch_size,
    shuffle=False
)

加载输出：

Found 8005 files belonging to 2 classes.
Using 6404 files for training.
Found 8005 files belonging to 2 classes.
Using 1601 files for validation.
Found 2023 files belonging to 2 classes.

用户尝试获取类别的代码：

class_indices = {}
for key in train_ds.class_indices.keys():
    class_indices[train_ds.class_indices[key]] = key

print(class_indices)

解决方案 #

方法1：直接使用Dataset对象的class_names属性 #

image_dataset_from_directory返回的_BatchDataset对象提供了class_names属性，它是一个按字母顺序排序的类别名称列表，列表的索引就是该类别对应的标签值。

如果需要生成{索引:类别名称}的字典，可以这样写：

class_indices = {idx: name for idx, name in enumerate(train_ds.class_names)}
print(class_indices)

方法2：手动从数据集目录读取类别 #

因为image_dataset_from_directory是按目录名的字母顺序识别类别的，也可以直接读取训练目录下的子目录名来获取类别：

import os

# 按字母顺序排序子目录名，确保和数据集的类别顺序一致
class_names = sorted(os.listdir(train_path))
class_indices = {idx: name for idx, name in enumerate(class_names)}
print(class_indices)

注意：需要确保train_path下只有类别对应的子目录，没有其他无关文件或目录。

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