How can I create a TensorFlow image dataset labelled by filenames from multiple directories?-CodePudding

I found a similar question but it doesn't quite fit this scenario.

I have some images of multiple classes in the same directory that I want to label according to their filename. I also have other directories with more images that I want to include in the same dataset.

Here's an example of the folder structure:

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

Is there any efficient way to create a dataset from this structure? Since I usually use image_dataset_from_directory() but it doesn't work in this case without moving images around.

CodePudding user response：

You can store all the file paths in a list and create a tensorflow dataset with it.

A simple way to get all the file paths is

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

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

Full example I created some sample images here on 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]

Creating dataset

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)

CodePudding user response：

There is no method that lets you do this automatically like image_dataset_from_directory (that I know of at least). However creating the dataset on your own is pretty easy.

You first have to iterate through your directory searching for the images. You save the path of each image you encounter in a list paths. For each image you check the class, either dog or cat in this case, and you create a dictionary where you map each class name to a different int, the label of that class. For each image you save aside the label of it in a list labels.

import glob

paths = []
labels = []
labels_dict = {}
index = 0

for filepath in glob.iglob('train/*/*.jpg'):
  paths.append(filepath)
  class_name = filepath.split('_')[1].split("/")[-1]
  if class_name not in labels_dict:
    labels_dict[class_name] = index
    index  = 1
  labels.append(labels_dict[class_name])

print(labels)
print(paths)
print(labels_dict)

Outputs:

# labels:
[0, 1, 1, 1, 1, 0]

# paths
['train/dir_1/dog_1.jpg', 'train/dir_1/cat_2.jpg', 'train/dir_1/cat_1.jpg', 'train/dir_2/cat_1.jpg', 'train/dir_2/cat_2.jpg', 'train/dir_2/dog_1.jpg']

# labels_dict
{'dog': 0, 'cat': 1}

Now that you have prepared your paths and labels, you can proceed as from this answer:

import tensorflow as tf
paths = tf.constant(paths)
labels = tf.constant(labels)

# create a dataset returning slices of `paths`
dataset = tf.data.Dataset.from_tensor_slices((paths, labels))

# parse every image in the dataset using `map`
def _parse_function(filename, label):
    image_string = tf.io.read_file(filename)
    image_decoded = tf.image.decode_jpeg(image_string, channels=3)
    image = tf.cast(image_decoded, tf.float32)
    return image, label

dataset = dataset.map(_parse_function)