如何从含图片路径与标签的Pandas DataFrame创建TensorFlow图像数据集

核心流程 #

通过tf.data.Dataset从DataFrame的图片路径和标签构建数据集，配合自定义预处理函数完成图像加载、格式转换与优化，最终得到可用于训练/验证的TensorFlow数据集。

代码实现步骤 #

1. 导入依赖库

import pandas as pd
import tensorflow as tf
from tensorflow.keras.layers import Rescaling

2. 加载DataFrame（已完成的话可跳过）

# 示例加载方式，替换为你的.csv路径
# df = pd.read_csv('image_labels.csv')

3. 初始化基础数据集

从DataFrame的img_path和label列生成TensorFlow数据集：

# 提取路径与标签，转为张量数据集
dataset = tf.data.Dataset.from_tensor_slices((df['img_path'].values, df['label'].values))

4. 定义图像加载与预处理函数

实现图片读取、解码、尺寸调整与归一化：

def load_and_preprocess_image(img_path, label):
    # 读取本地图片文件
    img_raw = tf.io.read_file(img_path)
    # 解码为RGB格式（灰度图用tf.image.decode_grayscale）
    img = tf.image.decode_jpeg(img_raw, channels=3)
    # 统一图像尺寸（根据你的模型需求修改）
    img = tf.image.resize(img, (224, 224))
    # 归一化像素值至[0,1]区间
    img = Rescaling(1./255)(img)
    return img, label

5. 预处理并优化数据集

通过map应用预处理，再进行打乱、批处理与预取操作提升效率：

# 并行应用预处理函数
dataset = dataset.map(load_and_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)

# 打乱数据集（训练阶段建议开启）
dataset = dataset.shuffle(buffer_size=len(df))

# 设置批次大小
batch_size = 32
dataset = dataset.batch(batch_size)

# 预取数据，加速后续训练
dataset = dataset.prefetch(tf.data.AUTOTUNE)

6. 验证数据集（可选）

取出一个批次验证数据是否正常：

for imgs_batch, labels_batch in dataset.take(1):
    print(f"单批次图像形状: {imgs_batch.shape}")
    print(f"前5个标签示例: {labels_batch[:5].numpy()}")

额外注意事项 #

• 标签编码：如果label列是字符串分类名称，需先转为整数编码：

  from sklearn.preprocessing import LabelEncoder
  
  le = LabelEncoder()
  df['label'] = le.fit_transform(df['label'])
  

• 数据增强：训练时可在预处理函数中添加图像增强操作（如翻转、旋转）：

  def load_and_preprocess_image(img_path, label):
      # 原有步骤...
      # 添加随机水平翻转
      img = tf.image.random_flip_left_right(img)
      return img, label
  

• 路径问题：确保img_path列的路径是绝对路径，或当前工作目录与相对路径匹配。

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