Title here
Summary here
ML Projects: Image Classification with CNNs
Project 2: Image Classification with CNNs
Introduction
Image classification is a fundamental task in computer vision, widely used in applications like medical imaging, autonomous driving, and facial recognition. In this project, we will develop a deep learning model using Convolutional Neural Networks (CNNs) to classify images.
Step 1: Dataset Overview
For this project, we will use the CIFAR-10 dataset, which contains 60,000 images across 10 categories (e.g., airplanes, cars, birds, cats, etc.). This dataset is publicly available in TensorFlow and PyTorch.
Accessing the Dataset
import tensorflow as tf
from tensorflow.keras.datasets import cifar10
# Load CIFAR-10 dataset
(X_train, y_train), (X_test, y_test) = cifar10.load_data()
# Normalize pixel values
X_train, X_test = X_train / 255.0, X_test / 255.0Step 2: Building a CNN Model
We will create a simple CNN architecture using TensorFlow and Keras.
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
# Define the CNN model
model = Sequential([
Conv2D(32, (3,3), activation='relu', input_shape=(32, 32, 3)),
MaxPooling2D(2,2),
Conv2D(64, (3,3), activation='relu'),
MaxPooling2D(2,2),
Flatten(),
Dense(128, activation='relu'),
Dropout(0.5),
Dense(10, activation='softmax')
])
# Compile the model
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])Step 3: Data Augmentation to Improve Performance
Data augmentation helps improve generalization by artificially expanding the training dataset.
from tensorflow.keras.preprocessing.image import ImageDataGenerator
datagen = ImageDataGenerator(
rotation_range=20,
width_shift_range=0.2,
height_shift_range=0.2,
horizontal_flip=True
)
# Fit the generator
datagen.fit(X_train)Step 4: Training the CNN Model
We will train our CNN model using augmented data.
# Train the model
history = model.fit(datagen.flow(X_train, y_train, batch_size=32),
validation_data=(X_test, y_test),
epochs=20)Step 5: Fine-Tuning a Pre-Trained Model (VGG16)
Instead of training from scratch, we can fine-tune a pre-trained model like VGG16.
from tensorflow.keras.applications import VGG16
from tensorflow.keras.models import Model
from tensorflow.keras.layers import GlobalAveragePooling2D
# Load VGG16 without the top classification layer
base_model = VGG16(weights='imagenet', include_top=False, input_shape=(32, 32, 3))
# Freeze pre-trained layers
for layer in base_model.layers:
layer.trainable = False
# Add new layers
x = GlobalAveragePooling2D()(base_model.output)
x = Dense(128, activation='relu')(x)
x = Dropout(0.5)(x)
out = Dense(10, activation='softmax')(x)
# Define new model
model = Model(inputs=base_model.input, outputs=out)
# Compile model
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])Step 6: Deploying the Model with Flask
We can deploy our trained model as an API using Flask.
Create a Flask API
from flask import Flask, request, jsonify
import pickle
import numpy as np
from tensorflow.keras.models import load_model
# Save the trained model
model.save('cnn_model.h5')
# Load the model
model = load_model('cnn_model.h5')
app = Flask(__name__)
@app.route('/predict', methods=['POST'])
def predict():
data = request.get_json()
image = np.array(data['image']).reshape((1, 32, 32, 3))
prediction = model.predict(image)
return jsonify({'class': int(np.argmax(prediction))})
if __name__ == '__main__':
app.run(debug=True)Testing the API
Save the script as app.py and run it. Then, test it using Postman or cURL.
curl -X POST http://127.0.0.1:5000/predict -H "Content-Type: application/json" -d '{"image": [pixel_values]}'Conclusion
In this project, we built a CNN model for image classification, applied data augmentation, fine-tuned VGG16, and deployed the model as an API. This project can be extended by using more complex architectures, optimizing hyperparameters, or deploying the model in a production environment.