Deep Learning
Keras
Image Preprocessing
Machine Learning
Data Augmentation

Keras Image Preprocessing

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

Introduction

Image preprocessing is the step where raw image files become tensors a neural network can learn from. In Keras, this stage is often the difference between a model that converges consistently and one that overfits or fails to train. Good preprocessing makes input shape, scale, and distribution predictable. It also increases data diversity through augmentation so the model sees more realistic variation during training.

Modern Keras pipelines usually rely on tf.data plus preprocessing layers, rather than older generator-only patterns. This approach is faster, easier to deploy, and keeps training and inference transformations aligned. The core goal is simple: make sure every pixel the model receives is in the expected format, and only apply random transforms where they improve generalization.

Core Sections

Build a clean input pipeline

Start by loading images with deterministic size and label handling.

python
1import tensorflow as tf
2
3IMG_SIZE = (224, 224)
4BATCH_SIZE = 32
5
6train_ds = tf.keras.utils.image_dataset_from_directory(
7    "data/train",
8    label_mode="int",
9    image_size=IMG_SIZE,
10    batch_size=BATCH_SIZE,
11    shuffle=True,
12    seed=42,
13)
14
15val_ds = tf.keras.utils.image_dataset_from_directory(
16    "data/val",
17    label_mode="int",
18    image_size=IMG_SIZE,
19    batch_size=BATCH_SIZE,
20    shuffle=False,
21)

This ensures all samples share shape and label type. Use a fixed seed for reproducibility when debugging.

Normalize and cache correctly

Neural networks train more stably when input values are normalized.

python
1normalizer = tf.keras.layers.Rescaling(1.0 / 255)
2
3AUTOTUNE = tf.data.AUTOTUNE
4train_ds = train_ds.map(lambda x, y: (normalizer(x), y), num_parallel_calls=AUTOTUNE)
5val_ds = val_ds.map(lambda x, y: (normalizer(x), y), num_parallel_calls=AUTOTUNE)
6
7train_ds = train_ds.cache().prefetch(AUTOTUNE)
8val_ds = val_ds.cache().prefetch(AUTOTUNE)

cache() improves throughput if your dataset fits memory or uses local SSD. prefetch() overlaps CPU preprocessing with GPU training.

Use augmentation layers for training only

Augmentation helps generalization by simulating variation in real-world inputs.

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
8inputs = tf.keras.Input(shape=(224, 224, 3))
9x = data_augmentation(inputs)
10x = tf.keras.applications.MobileNetV2(
11    include_top=False,
12    input_tensor=x,
13    pooling="avg",
14    weights="imagenet",
15).output
16outputs = tf.keras.layers.Dense(10, activation="softmax")(x)
17model = tf.keras.Model(inputs, outputs)

Because augmentation is inside the model graph, it runs automatically during training and is disabled during inference when configured in training mode logic.

Match preprocessing to pretrained backbones

If you use transfer learning, apply the exact preprocessing function required by that architecture.

python
1from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
2
3train_ds = train_ds.map(
4    lambda x, y: (preprocess_input(x * 255.0), y),
5    num_parallel_calls=AUTOTUNE,
6)

Different backbones expect different value ranges or color normalization. Mismatched preprocessing can silently degrade accuracy.

Handle class imbalance and label quality

Preprocessing is not just pixel transforms. You also need consistent labels and sampling strategy.

python
1class_weight = {
2    0: 1.0,
3    1: 2.3,
4    2: 1.8,
5}
6
7model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
8model.fit(train_ds, validation_data=val_ds, class_weight=class_weight, epochs=10)

If one class dominates, augmentation alone may not fix bias. Combine augmentation with weighted loss, better labels, or balanced sampling.

Common Pitfalls

  • Applying random augmentation to validation or test datasets, which makes evaluation unstable and not comparable across runs.
  • Mixing resizing methods between training and inference, causing subtle distribution shifts and lower production accuracy.
  • Forgetting backbone-specific preprocessing when using pretrained models, leading to poor transfer performance.
  • Caching a very large dataset in memory without checking limits, which can cause out-of-memory crashes.
  • Assuming augmentation quality compensates for noisy labels; mislabeled data still caps model performance.

Summary

Keras image preprocessing should be designed as a reproducible data pipeline, not a loose collection of transforms. Start with deterministic loading and normalization, add augmentation only where it helps training, and align preprocessing with your backbone model requirements. Then optimize throughput with cache and prefetch, while monitoring label quality and class balance. When these fundamentals are in place, training becomes faster, metrics become more stable, and the model is more likely to generalize to real images.

Treat preprocessing as part of the model contract and version it with your training code. If deployment preprocessing differs even slightly from training preprocessing, real-world accuracy can drop quickly. Keeping one shared pipeline definition across experiments, evaluation, and serving prevents that class of regression.


Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

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

Start the free 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

All Rights Reserved.