Introduction

image_dataset_from_directory gives you a tf.data.Dataset, not a static list of images. That means exploration and modification happen through dataset operations such as take, map, filter, unbatch, batch, cache, and prefetch, rather than by mutating some in-memory collection directly.

Inspect What the Dataset Actually Contains

Start by treating the dataset as a stream of batches. Check its metadata and inspect a batch before changing anything.

1import tensorflow as tf
2
3train_ds = tf.keras.utils.image_dataset_from_directory(
4    "data/train",
5    labels="inferred",
6    label_mode="int",
7    image_size=(224, 224),
8    batch_size=32,
9    shuffle=True,
10    seed=42,
11)
12
13print(train_ds.class_names)
14print(train_ds.element_spec)
15
16for images, labels in train_ds.take(1):
17    print(images.shape)
18    print(labels.shape)

element_spec tells you the tensor structure, while class_names tells you how folder names map to label indices.

Visualize Samples Early

A quick visualization often catches label mistakes, broken images, and unexpected resizing.

1import matplotlib.pyplot as plt
2
3for images, labels in train_ds.take(1):
4    plt.figure(figsize=(8, 8))
5    for i in range(9):
6        ax = plt.subplot(3, 3, i + 1)
7        plt.imshow(images[i].numpy().astype("uint8"))
8        plt.title(train_ds.class_names[int(labels[i])])
9        plt.axis("off")
10    plt.show()

This is often more informative than inspecting tensor shapes alone.

Modify the Dataset with map

Most transformations belong in map. This is where you normalize images, cast dtypes, change labels, or attach augmentation.

1def normalize(images, labels):
2    images = tf.cast(images, tf.float32) / 255.0
3    return images, labels
4
5train_ds = train_ds.map(normalize, num_parallel_calls=tf.data.AUTOTUNE)

This keeps preprocessing tied directly to the dataset pipeline instead of scattering it around the training loop.

Add Augmentation Carefully

If you want to modify the images for training, wrap augmentation into the dataset pipeline. Keep augmentation only on the training dataset, not validation or test data.

1augment = tf.keras.Sequential([
2    tf.keras.layers.RandomFlip("horizontal"),
3    tf.keras.layers.RandomRotation(0.1),
4    tf.keras.layers.RandomZoom(0.1),
5])
6
7
8def augment_batch(images, labels):
9    return augment(images, training=True), labels
10
11train_ds = train_ds.map(augment_batch, num_parallel_calls=tf.data.AUTOTUNE)

This changes the images produced by the dataset without touching the original files.

Filter or Remap Data

Because the result is a tf.data.Dataset, you can filter classes or remap labels using dataset operators.

train_ds = train_ds.unbatch()
train_ds = train_ds.filter(lambda x, y: tf.logical_or(tf.equal(y, 0), tf.equal(y, 1)))
train_ds = train_ds.batch(32)

Or remap labels:

1def remap(images, labels):
2    labels = tf.where(tf.equal(labels, 2), 1, labels)
3    return images, labels
4
5train_ds = train_ds.map(remap, num_parallel_calls=tf.data.AUTOTUNE)

This is useful when merging classes, dropping classes, or adapting a dataset to a different output structure.

Be Deliberate About Unbatching

unbatch is powerful, but it changes the dataset shape and affects ordering and performance. Use it only when the transformation genuinely needs per-example handling.

After unbatching and modifying, you usually need to batch again before training. That means it is easy to accidentally lose the batching and shuffling behavior you expected if you are not explicit.

Improve Throughput with Cache and Prefetch

Once the dataset pipeline is correct, make it efficient.

train_ds = train_ds.cache().prefetch(tf.data.AUTOTUNE)

For datasets that do not fit comfortably in memory, cache to disk or skip caching altogether. prefetch is usually a good default because it overlaps input work with model execution.

Validation Pipelines Should Differ from Training Pipelines

Create validation data with the same resizing and normalization, but without random augmentation.

1val_ds = tf.keras.utils.image_dataset_from_directory(
2    "data/val",
3    labels="inferred",
4    label_mode="int",
5    image_size=(224, 224),
6    batch_size=32,
7    shuffle=False,
8)
9
10val_ds = val_ds.map(normalize, num_parallel_calls=tf.data.AUTOTUNE)
11val_ds = val_ds.cache().prefetch(tf.data.AUTOTUNE)

Keeping the pipelines aligned except for augmentation prevents evaluation drift.

Common Pitfalls

A common mistake is treating the dataset like a mutable Python list. tf.data.Dataset is transformed functionally, not edited in place item by item.

Another mistake is applying augmentation to validation data, which makes metrics noisy and misleading.

Developers also often unbatch and rebatch without thinking about how that changes shuffle behavior and performance.

Finally, do not skip the initial inspection step. Many training problems come from dataset structure errors that were visible before the first epoch even started.

Summary

• 'image_dataset_from_directory returns a tf.data.Dataset, so exploration happens by iterating and inspection.'
• Use take, class_names, and element_spec to understand the dataset first.
• Use map, filter, unbatch, and batch to modify the pipeline.
• Keep augmentation only on the training path.
• Add cache and prefetch after correctness so the input pipeline stays efficient.