使用image_dataset_from_directory时如何获取批量图像文件名？

问题原因 #

image_dataset_from_directory生成的tf.data.Dataset默认仅输出图像张量与标签，且经过take、prefetch等转换操作后，原Dataset的file_paths属性会失效——因为这些操作会返回新的Dataset实例，不会继承原实例的属性。

解决方案 #

方案1：手动构建包含文件名的数据集（推荐，支持shuffle场景）

自己遍历文件路径，构建包含图像、标签、文件名的数据集，确保每个批量能直接获取对应文件名：

import tensorflow as tf
import pathlib

# 配置参数
TRAINING_DIR = "你的训练目录路径"
class_names = ['01', '02', '03']
Batch_Size = 32
Img_Size = (256, 256)

# 1. 遍历获取所有图像文件路径
data_dir = pathlib.Path(TRAINING_DIR)
# 根据你的图像格式调整后缀，比如*.jpg
file_paths = [str(path) for path in data_dir.glob('*/*.png')]

# 2. 生成对应分类标签（转为独热编码）
labels = [class_names.index(pathlib.Path(fp).parent.name) for fp in file_paths]
labels = tf.keras.utils.to_categorical(labels, num_classes=len(class_names))

# 3. 定义图像加载函数，返回图像、标签、文件名
def load_image(file_path, label):
    img = tf.io.read_file(file_path)
    # 灰度图设置channels=1，若为RGB则设为3；格式对应decode_jpg/decode_png
    img = tf.image.decode_png(img, channels=1)
    img = tf.image.resize(img, Img_Size, interpolation='bilinear')
    img = tf.cast(img, tf.float32) / 255.0  # 可选：归一化到[0,1]
    return img, label, file_path

# 4. 构建并优化数据集
trainData = tf.data.Dataset.from_tensor_slices((file_paths, labels))
trainData = trainData.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
trainData = trainData.batch(Batch_Size).prefetch(buffer_size=tf.data.AUTOTUNE)

# 5. 获取批量数据（包含文件名）
batch_data = trainData.as_numpy_iterator().next()
X_train, Y_train, batch_file_names = batch_data

方案2：基于原Dataset的file_paths映射（仅适用于shuffle=False场景）

如果保持使用image_dataset_from_directory，可以先保存原Dataset的file_paths，再根据批量顺序映射文件名：

import tensorflow as tf
from tensorflow.keras.utils import image_dataset_from_directory

# 配置参数
TRAINING_DIR = "你的训练目录路径"
class_names = ['01', '02', '03']
Batch_Size = 32
Img_Size = (256, 256)

# 1. 先构建原始数据集，暂不做take/prefetch
original_trainData = image_dataset_from_directory(
                                directory=TRAINING_DIR,
                                labels='inferred',
                                label_mode='categorical',
                                class_names=class_names,
                                color_mode='grayscale',
                                batch_size=Batch_Size,
                                image_size=Img_Size,
                                shuffle=False,
                                seed=42,
                                validation_split=None,
                                subset=None,
                                interpolation='bilinear',
                                follow_links=False,
                                crop_to_aspect_ratio=False) 

# 2. 保存所有文件路径
all_file_paths = original_trainData.file_paths

# 3. 执行后续优化操作
trainData = original_trainData.take(Batch_Size).prefetch(buffer_size=tf.data.AUTOTUNE)

# 4. 获取批量数据与对应文件名
X_train, Y_train = trainData.as_numpy_iterator().next()
# 因为shuffle=False，第一个batch对应前Batch_Size个路径
batch_file_paths = all_file_paths[:Batch_Size]

注意事项 #

• 方案1支持开启shuffle=True的场景，每个批量的文件名会和图像、标签一一对应，适合需要频繁关联文件名的任务；
• 方案2仅在shuffle=False时有效，若后续开启打乱，文件名与批量数据的顺序会错位；
• 针对千万级数据集，两种方案都能通过tf.data的并行机制保证加载效率，无需担心性能问题。

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