How do you load, label, and feed jpeg data into Tensorflow?

TensorFlow

JPEG

Data Loading

Image Processing

Machine Learning

How do you load, label, and feed jpeg data into Tensorflow?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

TensorFlow & Keras Computer Vision Machine Learning

Introduction

Loading JPEG images into TensorFlow involves reading image files, decoding them to tensors, resizing to a uniform shape, normalizing pixel values, associating labels, and batching into a tf.data.Dataset pipeline. The recommended approach is tf.keras.utils.image_dataset_from_directory, which handles all of these steps automatically when images are organized in label-named subdirectories. For custom loading, use tf.data.Dataset with tf.io.read_file and tf.io.decode_jpeg.

Method 1: image_dataset_from_directory (Recommended)

Organize images in directories named after their labels:

1data/
2  train/
3    cats/
4      cat001.jpg
5      cat002.jpg
6    dogs/
7      dog001.jpg
8      dog002.jpg
9  validation/
10    cats/
11      cat101.jpg
12    dogs/
13      dog101.jpg

python

1import tensorflow as tf
2
3# Load training data
4train_ds = tf.keras.utils.image_dataset_from_directory(
5    'data/train',
6    image_size=(224, 224),    # Resize all images
7    batch_size=32,
8    label_mode='int',         # Integer labels (0=cats, 1=dogs)
9    shuffle=True,
10    seed=42
11)
12
13# Load validation data
14val_ds = tf.keras.utils.image_dataset_from_directory(
15    'data/validation',
16    image_size=(224, 224),
17    batch_size=32,
18    label_mode='int'
19)
20
21# Check the dataset
22class_names = train_ds.class_names
23print(class_names)  # ['cats', 'dogs']
24
25for images, labels in train_ds.take(1):
26    print(f"Batch shape: {images.shape}")   # (32, 224, 224, 3)
27    print(f"Labels shape: {labels.shape}")  # (32,)
28    print(f"Pixel range: {images.numpy().min()} to {images.numpy().max()}")  # 0.0 to 255.0

Normalize and Optimize the Pipeline

python

1# Normalize pixel values to [0, 1]
2normalization = tf.keras.layers.Rescaling(1.0 / 255)
3train_ds = train_ds.map(lambda x, y: (normalization(x), y))
4val_ds = val_ds.map(lambda x, y: (normalization(x), y))
5
6# Optimize for performance
7AUTOTUNE = tf.data.AUTOTUNE
8train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
9val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

Method 2: Custom tf.data Pipeline

For images not organized in label directories:

python

1import tensorflow as tf
2import os
3
4# List of file paths and labels
5image_paths = ['path/to/img1.jpg', 'path/to/img2.jpg', 'path/to/img3.jpg']
6labels = [0, 1, 0]  # Corresponding labels
7
8def load_and_preprocess(path, label):
9    # Read the file
10    image = tf.io.read_file(path)
11
12    # Decode JPEG to tensor
13    image = tf.io.decode_jpeg(image, channels=3)  # Force 3 channels (RGB)
14
15    # Resize to uniform dimensions
16    image = tf.image.resize(image, [224, 224])
17
18    # Normalize to [0, 1]
19    image = image / 255.0
20
21    return image, label
22
23# Build dataset
24ds = tf.data.Dataset.from_tensor_slices((image_paths, labels))
25ds = ds.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
26ds = ds.shuffle(buffer_size=len(image_paths))
27ds = ds.batch(32)
28ds = ds.prefetch(tf.data.AUTOTUNE)
29
30# Use in training
31for images, labels in ds:
32    print(images.shape)  # (32, 224, 224, 3)

Method 3: Loading from CSV with Labels

python

1import pandas as pd
2import tensorflow as tf
3
4# CSV with columns: filepath, label
5df = pd.read_csv('labels.csv')
6# filepath,label
7# images/img001.jpg,cat
8# images/img002.jpg,dog
9
10# Encode string labels to integers
11label_to_int = {'cat': 0, 'dog': 1}
12df['label_int'] = df['label'].map(label_to_int)
13
14def load_image(path, label):
15    image = tf.io.read_file(path)
16    image = tf.io.decode_jpeg(image, channels=3)
17    image = tf.image.resize(image, [224, 224])
18    image = image / 255.0
19    return image, label
20
21ds = tf.data.Dataset.from_tensor_slices(
22    (df['filepath'].values, df['label_int'].values)
23)
24ds = ds.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
25ds = ds.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)

Data Augmentation

Apply random transformations to training images:

python

1data_augmentation = tf.keras.Sequential([
2    tf.keras.layers.RandomFlip("horizontal"),
3    tf.keras.layers.RandomRotation(0.1),
4    tf.keras.layers.RandomZoom(0.1),
5    tf.keras.layers.RandomContrast(0.1),
6])
7
8# Apply augmentation only to training data
9train_ds = train_ds.map(
10    lambda x, y: (data_augmentation(x, training=True), y),
11    num_parallel_calls=tf.data.AUTOTUNE
12)

Training a Model with the Dataset

python

1model = tf.keras.Sequential([
2    tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(224, 224, 3)),
3    tf.keras.layers.MaxPooling2D(),
4    tf.keras.layers.Conv2D(64, 3, activation='relu'),
5    tf.keras.layers.MaxPooling2D(),
6    tf.keras.layers.Conv2D(128, 3, activation='relu'),
7    tf.keras.layers.GlobalAveragePooling2D(),
8    tf.keras.layers.Dense(128, activation='relu'),
9    tf.keras.layers.Dense(2, activation='softmax')
10])
11
12model.compile(
13    optimizer='adam',
14    loss='sparse_categorical_crossentropy',
15    metrics=['accuracy']
16)
17
18history = model.fit(
19    train_ds,
20    validation_data=val_ds,
21    epochs=10
22)

Handling Large Datasets

For datasets that do not fit in memory:

python

1# TFRecord format — most efficient for large datasets
2def serialize_example(image_path, label):
3    image = tf.io.read_file(image_path)
4    feature = {
5        'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image.numpy()])),
6        'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))
7    }
8    example = tf.train.Example(features=tf.train.Features(feature=feature))
9    return example.SerializeToString()
10
11# Write TFRecords
12with tf.io.TFRecordWriter('data.tfrecord') as writer:
13    for path, label in zip(image_paths, labels):
14        writer.write(serialize_example(path, label))
15
16# Read TFRecords
17def parse_tfrecord(example_proto):
18    feature_desc = {
19        'image': tf.io.FixedLenFeature([], tf.string),
20        'label': tf.io.FixedLenFeature([], tf.int64)
21    }
22    parsed = tf.io.parse_single_example(example_proto, feature_desc)
23    image = tf.io.decode_jpeg(parsed['image'], channels=3)
24    image = tf.image.resize(image, [224, 224]) / 255.0
25    return image, parsed['label']
26
27ds = tf.data.TFRecordDataset('data.tfrecord')
28ds = ds.map(parse_tfrecord).batch(32).prefetch(tf.data.AUTOTUNE)

Common Pitfalls

Not specifying channels=3 in decode_jpeg: Some JPEG files are grayscale (1 channel) or CMYK (4 channels). Without channels=3, the tensor shape varies per image, causing batch errors. Always set channels=3 for RGB consistency.
Forgetting to normalize pixel values: Raw JPEG pixels range from 0 to 255. Neural networks train better with normalized inputs (0 to 1 or -1 to 1). Add image / 255.0 or use tf.keras.layers.Rescaling.
Images of different sizes without resizing: JPEGs come in various dimensions. Without tf.image.resize, tensors cannot be batched together because shapes differ. Always resize to a uniform (height, width).
Shuffling after batching: ds.batch(32).shuffle(100) shuffles batches, not individual samples. Place .shuffle() before .batch() for proper sample-level shuffling.
Loading all images into memory at once: Using tf.io.read_file inside Dataset.map loads images lazily. Avoid reading all images into a NumPy array first — that defeats the purpose of the tf.data pipeline and may cause out-of-memory errors.

Summary

Use tf.keras.utils.image_dataset_from_directory for the simplest loading when images are organized by label directories
Use tf.data.Dataset with tf.io.decode_jpeg for custom loading from file paths and labels
Always resize images to uniform dimensions and normalize pixel values before training
Add .cache().shuffle().prefetch(AUTOTUNE) for optimal pipeline performance
Use TFRecord format for large datasets that benefit from sequential read performance

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Basis Lab

No course covers this one yet

Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.

Build my course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

How do you load, label, and feed jpeg data into Tensorflow?

ML System Design practice on Codemia

Introduction

Method 1: image_dataset_from_directory (Recommended)

Normalize and Optimize the Pipeline

Method 2: Custom tf.data Pipeline

Method 3: Loading from CSV with Labels

Data Augmentation

Training a Model with the Dataset

Handling Large Datasets

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

No course covers this one yet

Track what you have practised

ML System Design practice on Codemia