TensorFlow Image Data Loading

Loading image data efficiently requires decoding image files, resizing them to a consistent shape, normalizing pixel values, and batching them for GPU processing. TensorFlow provides both high-level utilities that handle everything automatically and low-level tools for custom pipelines. This topic covers both approaches so you can choose what fits your project.

Method 1 — image_dataset_from_directory (Fastest Start)

If your images are organized into subfolders by class, this single function creates a complete tf.data pipeline automatically.

Expected folder structure:
data/
├── train/
│   ├── cats/
│   │   ├── cat001.jpg
│   │   └── cat002.jpg
│   └── dogs/
│       ├── dog001.jpg
│       └── dog002.jpg
└── val/
    ├── cats/
    └── dogs/

import tensorflow as tf

train_ds = tf.keras.utils.image_dataset_from_directory(
    'data/train/',
    image_size=(224, 224),   # Resize all images to this size
    batch_size=32,
    label_mode='int',        # Labels as integers (0, 1, 2...)
    shuffle=True,
    seed=42
)

val_ds = tf.keras.utils.image_dataset_from_directory(
    'data/val/',
    image_size=(224, 224),
    batch_size=32,
    label_mode='int',
    shuffle=False
)

# The class names are inferred from folder names
print(train_ds.class_names)  # ['cats', 'dogs']

Checking the Dataset Output

for images, labels in train_ds.take(1):
    print("Images batch shape:", images.shape)  # (32, 224, 224, 3)
    print("Labels batch shape:", labels.shape)  # (32,)
    print("Pixel value range:", images.numpy().min(), "–", images.numpy().max())
    # 0.0 – 255.0 (raw pixels — normalize before training)

Normalizing Pixel Values

Pixel values range from 0 to 255. Neural networks train best when inputs are small numbers close to zero. Normalize to [0, 1] range or [-1, 1] range depending on the model.

# Option 1: Rescaling layer (embedded in model — gets saved with it)
normalization_layer = tf.keras.layers.Rescaling(1./255)
train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y))

# Option 2: Divide in a preprocessing function
def normalize_images(images, labels):
    images = tf.cast(images, tf.float32) / 255.0
    return images, labels

train_ds = train_ds.map(normalize_images, num_parallel_calls=tf.data.AUTOTUNE)

# Option 3: Add Rescaling as first layer inside the model
model = tf.keras.Sequential([
    tf.keras.layers.Rescaling(1./255, input_shape=(224, 224, 3)),
    tf.keras.layers.Conv2D(32, 3, activation='relu'),
    ...
])

Adding Prefetch for Speed

AUTOTUNE = tf.data.AUTOTUNE

train_ds = (
    tf.keras.utils.image_dataset_from_directory('data/train/', image_size=(224,224), batch_size=32)
    .map(normalize_images, num_parallel_calls=AUTOTUNE)
    .cache()
    .shuffle(1000)
    .prefetch(AUTOTUNE)
)

val_ds = (
    tf.keras.utils.image_dataset_from_directory('data/val/', image_size=(224,224), batch_size=32)
    .map(normalize_images, num_parallel_calls=AUTOTUNE)
    .cache()
    .prefetch(AUTOTUNE)
)

Method 2 — Manual File Loading With tf.data

For more control over loading — custom folder structures, mixed data types, or metadata-based labels — build the pipeline manually.

import tensorflow as tf
import pathlib

data_dir = pathlib.Path('data/train/')

# Get all image file paths and their labels
image_paths = list(data_dir.glob('*/*.jpg'))
class_names = sorted([item.name for item in data_dir.glob('*') if item.is_dir()])
class_to_index = {name: i for i, name in enumerate(class_names)}

labels = [class_to_index[path.parent.name] for path in image_paths]

# Create path and label dataset
path_ds = tf.data.Dataset.from_tensor_slices(
    ([str(p) for p in image_paths], labels)
)

def load_and_preprocess(path, label):
    image = tf.io.read_file(path)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, [224, 224])
    image = tf.cast(image, tf.float32) / 255.0
    return image, label

dataset = (
    path_ds
    .map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
    .shuffle(len(image_paths))
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)

Loading Images With tf.keras Built-in Datasets

TensorFlow includes several ready-to-use image datasets for learning and benchmarking:

# CIFAR-10: 60,000 color images in 10 classes (32×32 pixels)
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()

# MNIST: 70,000 handwritten digit images (28×28 grayscale)
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# Fashion-MNIST: 70,000 clothing item images (28×28 grayscale)
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# Normalize
x_train = x_train.astype('float32') / 255.0
x_test  = x_test.astype('float32') / 255.0

# Add channel dimension for grayscale (if needed for Conv2D)
x_train = x_train[..., tf.newaxis]  # (60000, 28, 28, 1)

Handling PNG vs JPEG Formats

def load_image(path, label):
    raw = tf.io.read_file(path)

    # Detect format from file extension
    if tf.strings.regex_full_match(path, r'.*\.png'):
        image = tf.image.decode_png(raw, channels=3)
    else:
        image = tf.image.decode_jpeg(raw, channels=3)

    image = tf.image.resize(image, [224, 224])
    image = tf.cast(image, tf.float32) / 255.0
    return image, label

Diagram — Image Pipeline Flow

/data/cats/cat001.jpg  ─────────────────────────────┐
/data/cats/cat002.jpg  ─── list_files() ──────────  │
/data/dogs/dog001.jpg  ─────────────────────────────┘
         │
         ▼
[tf.io.read_file(path)]       → raw bytes
         │
         ▼
[tf.image.decode_jpeg()]      → uint8 tensor (H × W × 3)
         │
         ▼
[tf.image.resize([224,224])]  → (224 × 224 × 3)
         │
         ▼
[Cast to float32 / 255.0]     → values in [0, 1]
         │
         ▼
[batch(32) → prefetch()]      → (32, 224, 224, 3)
         │
         ▼
[CNN Model Input]

Inspecting Image Batches

import matplotlib.pyplot as plt

for images, labels in train_ds.take(1):
    plt.figure(figsize=(10, 10))
    for i in range(9):
        ax = plt.subplot(3, 3, i + 1)
        plt.imshow(images[i].numpy())
        plt.title(class_names[labels[i]])
        plt.axis('off')
    plt.show()

Efficient image loading is critical for CNN training speed. A pipeline that reads images slowly forces the GPU to wait between batches, multiplying training time unnecessarily. The next topic covers data augmentation — artificially expanding your dataset by creating modified versions of existing images, which dramatically improves model accuracy on small datasets.