如何在image_dataset_from_directory中使用自定义图像预处理函数？

问题根源 #

image_dataset_from_directory返回的是tf.data.Dataset，其中图像为TensorFlow张量格式，而你的preprocess_image函数基于NumPy和OpenCV实现，直接用.map()会因格式不兼容无法生效，需要将预处理逻辑适配到TensorFlow数据流管道中。

方法一：用tf.py_function包装原有预处理函数 #

这种方式可以完全保留你原有的预处理逻辑，仅需少量适配：

1. 定义Tensor兼容的包装函数，实现张量与NumPy数组的转换：

import tensorflow as tf

def tf_preprocess(image, label):
    # 调用原预处理函数，指定输出类型为uint8
    processed_img = tf.py_function(
        func=preprocess_image,
        inp=[image],
        Tout=tf.uint8
    )
    # 强制固定输出形状，避免数据流管道报错
    processed_img = tf.reshape(processed_img, (IMG_WIDTH, IMG_HEIGHT, 3))
    # 转换为模型常用的float32格式并归一化（可根据模型需求调整）
    processed_img = tf.cast(processed_img, tf.float32) / 255.0
    return processed_img, label

2. 生成数据集并应用预处理：

# 生成原始数据集
train_ds = tf.keras.utils.image_dataset_from_directory(
    directory=os.path.join(data_dir, 'train_images'),
    batch_size=BATCH_SIZE,
    image_size=(IMG_WIDTH, IMG_HEIGHT),
    seed=123,
    label_mode='categorical'
)

# 应用预处理，开启CPU并行加速
train_ds = train_ds.map(tf_preprocess, num_parallel_calls=tf.data.AUTOTUNE)

# 可选：添加缓存、预取优化，提升训练速度
train_ds = train_ds.cache().prefetch(tf.data.AUTOTUNE)

方法二：使用image_dataset_from_dataframe（贴合原流程） #

如果仍想基于DataFrame管理数据（和原flow_from_dataframe逻辑一致），可以用tf.keras.utils.image_dataset_from_dataframe替代，再配合上述预处理函数：

train_ds = tf.keras.utils.image_dataset_from_dataframe(
    dataframe=df,
    x_col='id_code',
    y_col='diagnosis',
    directory=os.path.join(data_dir, 'train_images'),
    batch_size=BATCH_SIZE,
    image_size=(IMG_WIDTH, IMG_HEIGHT),
    seed=123,
    label_mode='categorical'
)

# 应用预处理
train_ds = train_ds.map(tf_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
train_ds = train_ds.cache().prefetch(tf.data.AUTOTUNE)

关键注意事项 #

• 图像格式校验：image_dataset_from_directory/image_dataset_from_dataframe默认用PIL读取图像，输出为RGB格式。原preprocess_image中的cv2.cvtColor(image, cv2.COLOR_BGR2RGB)会导致颜色反转，若无需格式转换可删除该行。
• 形状一致性：预处理后的图像必须保持固定形状，tf.reshape是为了避免数据流因形状变化报错。
• 性能优化：num_parallel_calls=tf.data.AUTOTUNE会自动利用CPU并行处理，大幅提升预处理效率。

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