Image Classifier with TensorFlow
This guide demonstrates how to build and train a Convolutional Neural Network (CNN) to classify images using TensorFlow and Keras. We will use the CIFAR-10 dataset, which consists of 60,000 32x32 color images in 10 classes.
Modules Used:
tensorflow: An open-source platform for machine learning.numpy: For numerical operations.- matplotlib: To visualize the data and training results.
Installation
pip install tensorflow numpy matplotlibThe Code
Save this as image_classifier.py.
import tensorflow as tf
from tensorflow.keras import datasets, layers, models
import matplotlib.pyplot as plt
import numpy as np
def run_classifier():
# 1. Load and Preprocess Data
print("Loading CIFAR-10 dataset...")
(train_images, train_labels), (test_images, test_labels) = datasets.cifar10.load_data()
# Normalize pixel values to be between 0 and 1
train_images, test_images = train_images / 255.0, test_images / 255.0
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer',
'dog', 'frog', 'horse', 'ship', 'truck']
# 2. Build the CNN Model
print("Building model...")
model = models.Sequential()
# Convolutional layers to extract features
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
# Dense layers for classification
model.add(layers.Flatten())
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(10)) # 10 output classes
model.summary()
# 3. Compile the Model
model.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])
# 4. Train the Model
print("Training model (this may take a while)...")
history = model.fit(train_images, train_labels, epochs=10,
validation_data=(test_images, test_labels))
# 5. Evaluate
print("\nEvaluating model...")
test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)
print(f"\nTest accuracy: {test_acc:.4f}")
# 6. Plot Training History
plt.plot(history.history['accuracy'], label='accuracy')
plt.plot(history.history['val_accuracy'], label = 'val_accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.ylim([0.5, 1])
plt.legend(loc='lower right')
plt.title("Training History")
plt.show()
if __name__ == "__main__":
run_classifier()Usage
python image_classifier.pyNote: Training a CNN on a CPU can be slow. If you have a compatible NVIDIA GPU, ensure you have installed the necessary CUDA drivers and the tensorflow[and-cuda] package for faster training.