需numpy/cv2预处理时从目录加载图像的最优实现方案

核心实现思路 #

和你熟悉的PyTorch自定义Dataset + DataLoader逻辑完全对齐，放弃受限制的image_dataset_from_directoryAPI，完全自定义数据流水线即可同时解决两个问题，灵活度和PyTorch实现一致。

具体实现步骤 #

• 步骤1：生成文件路径与标签的对应关系
  先扫描数据集目录，把所有图片路径、对应标签分别存为两个列表，这一步可以完全自定义拆分规则：直接调用sklearn.model_selection.KFold生成K折的训练/验证索引，拆分路径、标签列表即可，完全不受keras自带拆分规则的限制，也支持自定义样本过滤、不均衡采样等额外逻辑，和你在PyTorch中做的操作完全一致。
• 步骤2：自定义numpy/cv2预处理逻辑
  用tf.py_function包装你的预处理逻辑，不需要手动取消分批、转换numpy数组，tf会自动处理张量和numpy数组的转换，不会出现之前map操作的兼容bug：

  import cv2
  import numpy as np
  import tensorflow as tf
  
  def preprocess_fn(file_path, label):
      def _preprocess(file_path, label):
          # 此处直接写cv2/numpy的预处理代码即可
          img = cv2.imread(file_path.numpy().decode())
          img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
          # 自定义预处理操作：裁剪、滤波、数据增强等
          img = cv2.resize(img, (224, 224)) / 255.0
          return img.astype(np.float32), label.numpy()
      # 包装为tf可识别的函数，指定输出数据类型
      img, label = tf.py_function(_preprocess, inp=[file_path, label], Tout=[tf.float32, tf.int32])
      img.set_shape((224, 224, 3))
      label.set_shape(())
      return img, label
  

• 步骤3：组装tf.data流水线（对应PyTorch的DataLoader）
  直接用拆分好的路径、标签列表生成数据集，组装加载逻辑即可：

  # 单折加载示例，K折训练时每次替换对应的train_paths/val_paths和对应标签即可
  batch_size = 32
  train_ds = tf.data.Dataset.from_tensor_slices((train_paths, train_labels))
  train_ds = train_ds.shuffle(len(train_paths))
  train_ds = train_ds.map(preprocess_fn, num_parallel_calls=tf.data.AUTOTUNE)
  train_ds = train_ds.batch(batch_size)
  train_ds = train_ds.prefetch(tf.data.AUTOTUNE)
  
  val_ds = tf.data.Dataset.from_tensor_slices((val_paths, val_labels))
  val_ds = val_ds.map(preprocess_fn, num_parallel_calls=tf.data.AUTOTUNE)
  val_ds = val_ds.batch(batch_size)
  

  该流水线加载性能和原生keras API基本一致，同时支持所有自定义需求。

轻量替代方案（保留原有API的情况） #

如果不想修改现有image_dataset_from_directory的调用逻辑，也可以按以下方式改造：

• 拆分逻辑：提前把数据集目录按KFold规则拆分为多个子目录，每次训练指定对应子目录作为训练/验证路径即可
• 预处理逻辑：在map调用中直接使用上述tf.py_function包装的预处理方法，不需要手动取消分批

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