Skip to content

Image Classification with Python — Train a Model from Scratch

DodaTech 1 min read

In this tutorial, you'll learn about Image Classification with Python. We cover key concepts, practical examples, and best practices to help you understand and apply this topic effectively.

What You'll Learn

Build a convolutional neural network (CNN) to classify images into 10 categories, train it on CIFAR-10, and visualize the results.

Why It Matters

Image classification is the foundation of Computer Vision — used in medical imaging, autonomous vehicles, security cameras, and quality inspection.

Real-World Use

Classifying product photos in e-commerce, detecting defects in manufacturing, and identifying plants or animals in nature apps.

Setup

import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
import numpy as np

Step 1: Load CIFAR-10 Dataset

(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()

# Classes
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer',
               'dog', 'frog', 'horse', 'ship', 'truck']

# Normalize
x_train, x_test = x_train / 255.0, x_test / 255.0

print(f"Train: {x_train.shape}, Test: {x_test.shape}")

Step 2: Build the CNN

model = keras.Sequential([
    keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Conv2D(64, (3, 3), activation='relu'),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Conv2D(64, (3, 3), activation='relu'),
    keras.layers.Flatten(),
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(10)
])

model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

Step 3: Train

history = model.fit(
    x_train, y_train,
    epochs=10,
    validation_data=(x_test, y_test)
)

Step 4: Evaluate

test_loss, test_acc = model.evaluate(x_test, y_test, verbose=2)
print(f"\nTest accuracy: {test_acc:.4f}")

Step 5: Visualize Predictions

predictions = model.predict(x_test[:5])

for i in range(5):
    pred_class = class_names[np.argmax(predictions[i])]
    true_class = class_names[y_test[i][0]]
    print(f"Predicted: {pred_class}, Actual: {true_class}")

Data Augmentation (Better Accuracy)

augmentation = keras.Sequential([
    keras.layers.RandomFlip("horizontal"),
    keras.layers.RandomRotation(0.1),
    keras.layers.RandomZoom(0.1),
])

# Apply during training
model.fit(
    augmentation(x_train), y_train,
    epochs=10,
    validation_data=(x_test, y_test)
)

Data augmentation typically improves accuracy by 5-10%.

Built by the developers of DodaTech

Doda Browser, DodaZIP & Durga Antivirus Pro