如何将单目录下带CSV标签的图像数据集加载至TensorFlow？

因为tf.keras.utils.image_dataset_from_directory依赖按类别分文件夹的目录结构，而你的数据是单目录存储图像、CSV记录标签的形式，推荐用tf.data.Dataset自定义加载流程，灵活性更高。

具体实现步骤 #

1. 读取并预处理CSV标签 #

先把CSV中的文件名和文本标签转换成模型可识别的格式：

import pandas as pd
import tensorflow as tf

# 读取CSV文件
df = pd.read_csv('labels.csv')
# 将文本标签映射为数值（cat=0，dog=1）
label_mapping = {'cat': 0, 'dog': 1}
df['label'] = df['label'].map(label_mapping)

# 提取文件名和标签数组
filenames = df['filename'].values
labels = df['label'].values

2. 构建图像加载流水线 #

编写函数完成图像路径拼接、读取、解码与预处理，再通过tf.data.Dataset映射到数据集中：

def load_preprocess_img(filename, label):
    # 拼接图像完整路径（假设图像都在data目录下）
    img_full_path = tf.strings.join(['data/', filename])
    # 读取图像文件
    raw_img = tf.io.read_file(img_full_path)
    # 解码图像（如果是JPG格式，替换为tf.image.decode_jpeg）
    decoded_img = tf.image.decode_png(raw_img, channels=3)
    # 调整图像尺寸（根据你的模型需求修改，比如224x224）
    resized_img = tf.image.resize(decoded_img, [224, 224])
    # 归一化像素值到[0, 1]区间
    normalized_img = tf.cast(resized_img, tf.float32) / 255.0
    return normalized_img, label

# 从文件名和标签构建基础数据集
dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
# 映射加载预处理函数，开启多线程加速
dataset = dataset.map(load_preprocess_img, num_parallel_calls=tf.data.AUTOTUNE)

3. 优化数据集性能 #

添加打乱、分批、预取操作，提升训练时的效率：

# 打乱数据集（buffer_size建议设为样本总数，内存不足时可适当缩小）
dataset = dataset.shuffle(buffer_size=len(filenames))
# 设置批次大小
batch_size = 32
dataset = dataset.batch(batch_size)
# 预取数据，让GPU训练与CPU数据准备并行
dataset = dataset.prefetch(tf.data.AUTOTUNE)

4. 验证加载结果（可选） #

快速检查数据集是否正常：

# 取出第一个批次的图像和标签
for batch_imgs, batch_lbls in dataset.take(1):
    print(f"单批次图像形状: {batch_imgs.shape}")
    print(f"单批次标签形状: {batch_lbls.shape}")
    print(f"第一个图像的标签: {batch_lbls[0].numpy()}")

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