Multi-class image classification is a core computer vision task where you assign each image to one of several predefined categories. Keras, running on top of TensorFlow, gives you a high-level API that makes it straightforward to build, train, and evaluate convolutional neural networks (CNNs) for this purpose. This article walks you through the complete workflow, from loading data to evaluating your trained model.

Dataset Preparation

Keras provides a convenient utility for loading images from a directory structure where each subfolder name is the class label.

1import tensorflow as tf
2from tensorflow.keras.utils import image_dataset_from_directory
3
4train_ds = image_dataset_from_directory(
5    "data/train",
6    image_size=(128, 128),
7    batch_size=32,
8    label_mode="categorical",   # one-hot encoded labels
9)
10
11val_ds = image_dataset_from_directory(
12    "data/val",
13    image_size=(128, 128),
14    batch_size=32,
15    label_mode="categorical",
16)

The directory layout should look like this:

1data/
2  train/
3    cats/
4    dogs/
5    birds/
6  val/
7    cats/
8    dogs/
9    birds/

Setting label_mode="categorical" produces one-hot encoded labels, which is required when you use categorical_crossentropy as the loss function.

Building a CNN Model

A typical CNN for multi-class classification stacks convolutional and pooling layers to extract features, then uses dense layers to produce class probabilities. The final layer must have as many units as there are classes and use a softmax activation.

1from tensorflow.keras import layers, models
2
3num_classes = 3
4
5model = models.Sequential([
6    layers.Rescaling(1.0 / 255, input_shape=(128, 128, 3)),
7    layers.Conv2D(32, (3, 3), activation="relu"),
8    layers.MaxPooling2D((2, 2)),
9    layers.Conv2D(64, (3, 3), activation="relu"),
10    layers.MaxPooling2D((2, 2)),
11    layers.Conv2D(128, (3, 3), activation="relu"),
12    layers.MaxPooling2D((2, 2)),
13    layers.Flatten(),
14    layers.Dense(128, activation="relu"),
15    layers.Dropout(0.5),
16    layers.Dense(num_classes, activation="softmax"),
17])

The Rescaling layer normalizes pixel values from the 0-255 range down to 0-1, which helps the optimizer converge faster.

Compiling and Training

For multi-class classification, compile the model with categorical_crossentropy loss. Use an optimizer like Adam and track accuracy as the metric.

1model.compile(
2    optimizer="adam",
3    loss="categorical_crossentropy",
4    metrics=["accuracy"],
5)
6
7history = model.fit(
8    train_ds,
9    validation_data=val_ds,
10    epochs=20,
11)

If your labels are integers rather than one-hot vectors, use sparse_categorical_crossentropy instead and set label_mode="int" when loading the dataset.

Data Augmentation

Small datasets benefit greatly from data augmentation, which generates varied versions of each training image to reduce overfitting. In Keras you can add augmentation layers directly inside the model so that augmentation happens on the GPU during training.

1data_augmentation = tf.keras.Sequential([
2    layers.RandomFlip("horizontal"),
3    layers.RandomRotation(0.1),
4    layers.RandomZoom(0.1),
5])
6
7model = models.Sequential([
8    data_augmentation,
9    layers.Rescaling(1.0 / 255, input_shape=(128, 128, 3)),
10    layers.Conv2D(32, (3, 3), activation="relu"),
11    layers.MaxPooling2D((2, 2)),
12    layers.Conv2D(64, (3, 3), activation="relu"),
13    layers.MaxPooling2D((2, 2)),
14    layers.Conv2D(128, (3, 3), activation="relu"),
15    layers.MaxPooling2D((2, 2)),
16    layers.Flatten(),
17    layers.Dense(128, activation="relu"),
18    layers.Dropout(0.5),
19    layers.Dense(num_classes, activation="softmax"),
20])

These augmentation layers are only active during training. At inference time they pass images through unchanged.

Evaluation

After training, evaluate the model on a held-out test set to measure generalization performance.

1test_ds = image_dataset_from_directory(
2    "data/test",
3    image_size=(128, 128),
4    batch_size=32,
5    label_mode="categorical",
6)
7
8loss, accuracy = model.evaluate(test_ds)
9print(f"Test accuracy: {accuracy:.4f}")

For a more detailed breakdown, use a confusion matrix to see which classes the model confuses most often:

1import numpy as np
2from sklearn.metrics import classification_report
3
4y_true, y_pred = [], []
5for images, labels in test_ds:
6    preds = model.predict(images)
7    y_true.extend(np.argmax(labels.numpy(), axis=1))
8    y_pred.extend(np.argmax(preds, axis=1))
9
10print(classification_report(y_true, y_pred, target_names=["cats", "dogs", "birds"]))

Common Pitfalls

• Mismatching loss and label format. Using categorical_crossentropy with integer labels (or sparse_categorical_crossentropy with one-hot labels) causes shape errors or silent incorrect training.
• Forgetting to rescale pixel values. Raw pixel values in the 0-255 range produce very large activations. Always normalize inputs to 0-1 or use standardization.
• Setting the wrong number of output units. The final Dense layer must have exactly as many units as there are classes. A mismatch causes a shape error during training.
• Applying augmentation at test time. Data augmentation should only run during training. If you build augmentation outside the model and apply it to the test set, your evaluation metrics will be unreliable.
• Training on too few epochs or too many. Too few epochs means the model underfits; too many leads to overfitting. Monitor validation loss and use EarlyStopping to find the sweet spot automatically.

Summary

• Organize your images into subdirectories named by class, and use image_dataset_from_directory with label_mode="categorical" for one-hot labels.
• Build a CNN that ends with a Dense layer of size num_classes and softmax activation.
• Compile with categorical_crossentropy (or sparse_categorical_crossentropy for integer labels) and the Adam optimizer.
• Add augmentation layers inside the model to reduce overfitting on small datasets.
• Evaluate on a separate test set and use a confusion matrix or classification report to understand per-class performance.