创建自定义TF.DataSet图像数据集无法生成数据标签及配置路径问题

执行tfds new my_dataset只会生成数据集脚手架代码，不会自动识别你本地的Train目录结构、自动生成带标签的数据集，你需要手动修改模板中的路径配置、数据解析逻辑，再执行构建才能生成可用的数据集。

1. 修改核心模板文件my_dataset.py #

生成的3个脚本中，只有和数据集同名的my_dataset.py需要改动，按以下位置修改：

• 先在文件头部补充依赖导入：import os
• 修改_info()方法，定义数据集的特征结构：如果是图像分类任务，明确指定输入为图片格式，标签对应4个类别，类别名和Train下4个子文件夹名完全一致，示例：

def _info(self):
  return tfds.core.DatasetInfo(
      builder=self,
      features=tfds.features.FeaturesDict({
          'image': tfds.features.Image(),
          'label': tfds.features.ClassLabel(names=['cat', 'dog', 'bird', 'fish']), # 替换成你自己的4个子文件夹名
      }),
      supervised_keys=('image', 'label'),
  )

• 修改_split_generators()方法，删掉模板默认的远程下载逻辑，直接指定本地Train文件夹的绝对路径：

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
  # 替换成你本地Train文件夹的实际绝对路径
  train_data_root = "/home/xxx/project/Train"
  return {
      'train': self._generate_examples(train_data_root),
      # 如需验证集、测试集，按相同逻辑指定对应路径即可
  }

• 修改_generate_examples()方法，实现目录遍历、标签匹配逻辑，自动把不同子文件夹下的样本对应到正确标签：

def _generate_examples(self, data_path):
  class_names = self.info.features['label'].names
  for label_id, class_name in enumerate(class_names):
    class_folder = os.path.join(data_path, class_name)
    # 遍历当前类别文件夹下的所有样本文件
    for file_name in os.listdir(class_folder):
      # 过滤掉隐藏文件、非样本文件，按需调整后缀判断
      if file_name.startswith('.'):
        continue
      file_path = os.path.join(class_folder, file_name)
      # 每个样本返回唯一ID和对应的特征、标签
      yield f"{class_name}_{file_name}", {
          'image': file_path,
          'label': label_id
      }

2. 重新执行构建命令 #

修改完文件后，在my_dataset目录下重新执行构建命令：

tfds build

命令执行完成后，终端会输出构建成功提示，打包好的TFDS格式数据集会默认存到~/tensorflow_datasets/my_dataset/路径下。

3. 加载自定义数据集 #

构建无报错的话，直接用你预期的代码即可加载数据集：

import tensorflow_datasets as tfds
ds = tfds.load('my_dataset', split='train')

如果加载时提示找不到数据集，有两种解决方式：

• 把存放my_dataset.py的目录加入Python系统路径
• 在tfds.load中传入data_dir参数，显式指定~/tensorflow_datasets/路径

注意：如果你的数据不是图片格式，只需要对应修改_info()里的特征定义、_generate_examples()里的文件读取逻辑即可，整体流程不变。

内容的提问来源于stack exchange，提问作者Aneesh R P