Introduction

Before a neural network can learn from images, those images must be loaded from disk, decoded into numerical tensors, resized to a uniform shape, and paired with their labels. TensorFlow provides several APIs that handle this pipeline efficiently, from low-level file I/O to high-level directory loaders. This article shows you three progressively higher-level approaches so you can pick the one that matches your project's complexity.

Low-Level Loading with tf.io and tf.image

The most explicit way to load a JPEG is to read the raw bytes with tf.io.read_file and decode them with tf.image.decode_jpeg. This gives you full control over every preprocessing step.

1import tensorflow as tf
2
3def load_and_preprocess(file_path, label):
4    raw = tf.io.read_file(file_path)
5    image = tf.image.decode_jpeg(raw, channels=3)
6    image = tf.image.resize(image, [224, 224])
7    image = image / 255.0  # normalize to [0, 1]
8    return image, label

You then combine this function with a tf.data.Dataset built from your file paths and labels.

1file_paths = ["data/cats/001.jpg", "data/dogs/002.jpg"]
2labels = [0, 1]  # 0 = cat, 1 = dog
3
4dataset = tf.data.Dataset.from_tensor_slices((file_paths, labels))
5dataset = dataset.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
6dataset = dataset.shuffle(buffer_size=100)
7dataset = dataset.batch(32)
8dataset = dataset.prefetch(tf.data.AUTOTUNE)

num_parallel_calls=tf.data.AUTOTUNE lets TensorFlow decide how many images to decode in parallel, and prefetch ensures the next batch is being prepared while the GPU trains on the current one.

High-Level Loading with image_dataset_from_directory

If your images are organized in folders where each subfolder name is the class label, tf.keras.utils.image_dataset_from_directory handles loading, labeling, batching, and resizing in a single call.

1data/
2  cats/
3    001.jpg
4    002.jpg
5  dogs/
6    003.jpg
7    004.jpg

1import tensorflow as tf
2
3train_ds = tf.keras.utils.image_dataset_from_directory(
4    "data/",
5    image_size=(224, 224),
6    batch_size=32,
7    validation_split=0.2,
8    subset="training",
9    seed=42,
10)
11
12val_ds = tf.keras.utils.image_dataset_from_directory(
13    "data/",
14    image_size=(224, 224),
15    batch_size=32,
16    validation_split=0.2,
17    subset="validation",
18    seed=42,
19)
20
21# Normalize pixel values
22normalization = tf.keras.layers.Rescaling(1.0 / 255)
23train_ds = train_ds.map(lambda x, y: (normalization(x), y))
24val_ds = val_ds.map(lambda x, y: (normalization(x), y))

This utility infers labels from the folder names, splits training and validation sets, and returns a tf.data.Dataset ready for model.fit(). For most classification tasks this is the fastest way to get started.

Building a Custom Dataset with Labels from a CSV

Sometimes labels live in an external file (a CSV, a database, or a JSON manifest) rather than in the directory structure. In that case you can build the file-path-to-label mapping yourself and feed it into a tf.data.Dataset.

1import tensorflow as tf
2import pandas as pd
3
4df = pd.read_csv("labels.csv")  # columns: filename, label
5file_paths = df["filename"].tolist()
6labels = df["label"].tolist()
7
8# Build a class-name-to-integer mapping
9class_names = sorted(set(labels))
10class_to_idx = {name: i for i, name in enumerate(class_names)}
11int_labels = [class_to_idx[l] for l in labels]
12
13def load_and_preprocess(path, label):
14    raw = tf.io.read_file(path)
15    img = tf.image.decode_jpeg(raw, channels=3)
16    img = tf.image.resize(img, [224, 224])
17    img = img / 255.0
18    return img, label
19
20dataset = tf.data.Dataset.from_tensor_slices((file_paths, int_labels))
21dataset = (
22    dataset
23    .shuffle(len(file_paths))
24    .map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
25    .batch(32)
26    .prefetch(tf.data.AUTOTUNE)
27)

This pattern is flexible enough to handle multi-label classification, regression targets, or any labeling scheme that does not map neatly to a folder hierarchy.

Adding Data Augmentation

Once your dataset pipeline is in place, you can insert augmentation layers to improve generalization. TensorFlow's Keras preprocessing layers run on the GPU and integrate directly into the pipeline.

1augmentation = 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
7train_ds = train_ds.map(
8    lambda x, y: (augmentation(x, training=True), y),
9    num_parallel_calls=tf.data.AUTOTUNE,
10)

Apply augmentation only to the training set. Validation and test sets should use the original, unmodified images so your metrics reflect real-world performance.

Common Pitfalls

• Forgetting to normalize pixel values: JPEG pixels range from 0 to 255. Feeding raw values into a model with small initial weights produces enormous activations and unstable training. Always rescale to [0, 1] or [-1, 1].
• Mismatched image sizes: If you skip tf.image.resize, images of different dimensions will cause a shape error when TensorFlow tries to batch them. Every image in a batch must have the same height and width.
• Using shuffle with a buffer that is too small: tf.data.Dataset.shuffle(buffer_size=10) only shuffles within a window of 10 elements, which can leave the dataset nearly sorted. Set the buffer size to at least the number of samples in your dataset for a true shuffle.
• Applying augmentation to validation data: Random flips, rotations, and zooms should only be applied during training. Augmenting validation data gives you noisy metrics that do not reflect real inference performance.
• Decoding PNGs with decode_jpeg: tf.image.decode_jpeg will throw an error on PNG files. If your dataset contains mixed formats, use tf.image.decode_image instead, which auto-detects the format.

Summary

• Use tf.io.read_file and tf.image.decode_jpeg for full control over the loading and preprocessing pipeline.
• Use tf.keras.utils.image_dataset_from_directory when your images are organized in class-named subfolders for quick setup.
• Build a custom tf.data.Dataset from file paths and a CSV or manifest file when labels are stored externally.
• Chain .shuffle(), .batch(), and .prefetch(tf.data.AUTOTUNE) to keep the GPU fed efficiently during training.
• Normalize pixel values and apply data augmentation only to the training split, never to validation or test data.