TensorFlow Image Classification

Image classification assigns a label to an entire image. A photo of a cat gets the label "cat." A photo of a dog gets "dog." This task is the most common starting point for CNN projects and the benchmark used to measure progress in computer vision research. This topic builds a complete image classifier from data loading through training, evaluation, and prediction — applying every CNN concept covered so far.

Project: Classifying Flowers Into 5 Categories

TensorFlow provides a Flowers dataset with 3,670 images across five categories: daisy, dandelion, roses, sunflowers, and tulips. This is a realistic small-dataset classification problem.

Step 1 — Load and Explore the Dataset

import tensorflow as tf
import pathlib

# Download the dataset
dataset_url = 'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz'
data_dir = tf.keras.utils.get_file('flower_photos', origin=dataset_url,
                                    untar=True)
data_dir = pathlib.Path(data_dir)

# Count images
image_count = len(list(data_dir.glob('*/*.jpg')))
print(f"Total images: {image_count}")   # 3,670

# Class names from folder names
class_names = sorted([item.name for item in data_dir.glob('*')
                       if item.is_dir()])
print(class_names)   # ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']

Step 2 — Build the Data Pipeline

IMAGE_SIZE = (180, 180)
BATCH_SIZE = 32
AUTOTUNE = tf.data.AUTOTUNE

train_ds = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset='training',
    seed=42,
    image_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE
)

val_ds = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset='validation',
    seed=42,
    image_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE
)

# Optimize pipeline
train_ds = train_ds.cache().shuffle(1000).prefetch(AUTOTUNE)
val_ds   = val_ds.cache().prefetch(AUTOTUNE)

Step 3 — Build the CNN Model

import tensorflow as tf

NUM_CLASSES = 5

model = tf.keras.Sequential([
    # Augmentation (training only)
    tf.keras.layers.RandomFlip('horizontal', input_shape=(180, 180, 3)),
    tf.keras.layers.RandomRotation(0.1),
    tf.keras.layers.RandomZoom(0.1),

    # Normalization
    tf.keras.layers.Rescaling(1./255),

    # Block 1: 32 filters
    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),

    # Block 2: 64 filters
    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),

    # Block 3: 128 filters
    tf.keras.layers.Conv2D(128, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),

    # Block 4: 256 filters
    tf.keras.layers.Conv2D(256, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),

    # Classifier head
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')
])

model.summary()

Step 4 — Compile and Train

model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

history = model.fit(
    train_ds,
    epochs=30,
    validation_data=val_ds,
    callbacks=[
        tf.keras.callbacks.EarlyStopping(
            monitor='val_accuracy',
            patience=5,
            restore_best_weights=True
        )
    ]
)

print(f"Best validation accuracy: {max(history.history['val_accuracy']):.2%}")

Step 5 — Evaluate on the Full Validation Set

loss, accuracy = model.evaluate(val_ds)
print(f"Final validation loss:     {loss:.4f}")
print(f"Final validation accuracy: {accuracy:.2%}")

Step 6 — Predict on New Images

import numpy as np

def predict_image(model, image_path, class_names):
    # Load and preprocess a single image
    img = tf.keras.utils.load_img(image_path, target_size=(180, 180))
    img_array = tf.keras.utils.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)   # Add batch dimension

    predictions = model.predict(img_array, verbose=0)
    predicted_class = class_names[np.argmax(predictions[0])]
    confidence = np.max(predictions[0])

    print(f"Predicted: {predicted_class} ({confidence:.1%} confidence)")
    print(f"All scores: {dict(zip(class_names, predictions[0].round(3)))}")

predict_image(model, 'my_flower.jpg', class_names)

How the Model Sees the Image

Input Image: (180, 180, 3) — color photo of a flower

After Aug + Rescaling:  (180, 180, 3) — pixel values 0–1
After Conv2D(32)+Pool:  (90, 90, 32)  — 32 feature maps, half size
After Conv2D(64)+Pool:  (45, 45, 64)  — 64 feature maps, half size again
After Conv2D(128)+Pool: (22, 22, 128) — 128 maps, quarter size
After Conv2D(256)+Pool: (11, 11, 256) — 256 maps, eighth size
After Flatten:          (30976,)      — all features in a 1D vector
After Dense(128):       (128,)        — 128 abstract features
After Dense(5)+Softmax: (5,)          — one probability per class

[daisy:0.05, dandelion:0.03, roses:0.87, sunflowers:0.03, tulips:0.02]
Prediction → "roses"

Reading the Training History

history.history.keys()
# ['loss', 'accuracy', 'val_loss', 'val_accuracy']

# Identify overfitting
train_acc = history.history['accuracy']
val_acc   = history.history['val_accuracy']

for epoch, (t, v) in enumerate(zip(train_acc, val_acc), 1):
    gap = t - v
    status = 'OVERFITTING' if gap > 0.15 else 'OK'
    print(f"Epoch {epoch:02d}: train={t:.3f} val={v:.3f} gap={gap:.3f} {status}")

Common Classification Mistakes and Fixes

Symptom                               Fix
────────────────────────────────────────────────────────────────────
val_accuracy stuck at 1/N_CLASSES     Wrong loss function; use
(random guessing)                     sparse_categorical_crossentropy

Training accuracy high,               Add more Dropout, add augmentation,
val accuracy much lower               reduce model size, gather more data

Loss is NaN from epoch 1              Learning rate too high; try 1e-4 or 1e-5

Accuracy improves then                Use EarlyStopping + restore_best_weights
drops after many epochs

All predictions same class            Check class imbalance; add class_weight
────────────────────────────────────────────────────────────────────

Improving Accuracy Further

• Use transfer learning with MobileNetV2 or EfficientNet (covered in Topic 34–37)
• Add more augmentation: RandomBrightness, RandomContrast
• Use label smoothing: loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1)
• Try learning rate warmup and cosine decay schedule
• Collect more data for the weakest-performing class

Building a complete image classifier from raw images to accurate predictions is the core skill of computer vision engineering. The architecture pattern here — augmentation, rescaling, stacked Conv+Pool blocks, Dropout, Dense head — appears in virtually every practical image classification project. The next section shifts to Recurrent Neural Networks, which handle sequential data like text, audio, and time series.