tensorflow - 如何创建由多个目录中的文件名标记的 TensorFlow 图像数据集？

Question

我发现了一个类似的问题，但不太适合这种情况。

我在同一目录中有一些多个类的图像，我想根据它们的文件名来标记它们。我还有其他目录，其中包含更多图像，我想将它们包含在同一数据集中。

以下是文件夹结构的示例:

train/directory_1/cat_1.png
train/directory_1/cat_2.png
train/directory_1/dog_1.png
train/directory_1/dog_2.png
...
train/directory_2/cat_1.png
train/directory_2/cat_2.png
train/directory_2/dog_1.png
train/directory_2/dog_2.png

是否有任何有效的方法可以从该结构创建数据集？因为我通常使用 image_dataset_from_directory() 但在这种情况下如果不移动图像它就不起作用。

Answer 1

您可以将所有文件路径存储在列表中，并用它创建 tensorflow 数据集。

获取所有文件路径的简单方法是

import glob
paths = glob.glob("train/*/*")

#If you want only the .png files
paths = glob.glob("train/*/*.png")

完整示例我在 Colab 上创建了一些示例图像

paths = glob.glob("/content/*/*.jpeg")
paths
>>>
['/content/dir_1/cat_1.jpeg',
 '/content/dir_1/dog_1.jpeg',
 '/content/dir_2/dog_2.jpeg',
 '/content/dir_2/cat_2.jpeg']

paths = glob.glob("/content/*/*.jpeg")
label_dict = {'cat':0, 'dog':1}


def get_label(x):
  label = x.split('.')[0].rsplit('_')[-2].split('/')[-1]
  return label_dict[label]

labels = [get_label(i) for i in paths]
labels
>>>
[0, 1, 1, 0]

创建数据集

def preprocess(x,y):
  img = tf.io.read_file(x)
  img = tf.io.decode_png(img, channels=3)
  img = tf.cast(img,dtype=tf.float32)
  img = img / 255.
  
  return img,y

dataset = tf.data.Dataset.from_tensor_slices((paths,labels))
dataset = dataset.map(preprocess)
dataset = dataset.batch(2)

for x,y in dataset.take(1):
  print(x.shape, y)
>>>
(2, 192, 262, 3) tf.Tensor([0 1], shape=(2,), dtype=int32)