AI-powered image classification is a foundational task in computer vision with applications in areas like healthcare, retail, and security. In this tutorial, we’ll build a simple yet effective image classifier using TensorFlow and deploy it as a web service with FastAPI. This approach allows you to serve your model to any client via an easy-to-use REST API.
🔧 Step 1: Install Required Packages
pip install tensorflow fastapi uvicorn python-multipart pillow numpy
This installs TensorFlow for building the model, FastAPI and Uvicorn for serving the API, Pillow for image processing, and python-multipart to handle file uploads.
📁 Step 2: Project Structure
.
├── main.py
├── model/
│ └── saved_model/
├── requirements.txt
├── uploads/
🧠 Step 3: Train and Save a Simple Image Classifier
We’ll use TensorFlow’s built-in dataset (e.g., Fashion MNIST) for training a simple model and save it under model/saved_model/.
import tensorflow as tf
from tensorflow.keras import layers, models
# Load dataset
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()
# Normalize images
train_images = train_images / 255.0
test_images = test_images / 255.0
# Build model
model = models.Sequential([
layers.Flatten(input_shape=(28, 28)),
layers.Dense(128, activation='relu'),
layers.Dense(10, activation='softmax')
])
# Compile model
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
# Train model
model.fit(train_images, train_labels, epochs=5, validation_data=(test_images, test_labels))
# Save model
model.save("model/saved_model")
🔐 Step 4: Create main.py with FastAPI Serving Logic
This script loads the trained model and provides an endpoint to upload an image file and get the predicted class.
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import uvicorn
import numpy as np
from PIL import Image
import io
import tensorflow as tf
app = FastAPI()
# Load the saved TensorFlow model
model = tf.keras.models.load_model("model/saved_model")
# Class names for Fashion MNIST dataset
class_names = [
"T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
"Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]
def preprocess_image(image_bytes):
image = Image.open(io.BytesIO(image_bytes)).convert('L') # Convert to grayscale
image = image.resize((28, 28))
image_array = np.array(image) / 255.0 # Normalize
return image_array.reshape(1, 28, 28)
@app.post("/predict/")
async def predict_image(file: UploadFile = File(...)):
contents = await file.read()
img = preprocess_image(contents)
predictions = model.predict(img)
predicted_class = class_names[np.argmax(predictions)]
confidence = float(np.max(predictions))
return JSONResponse({
"filename": file.filename,
"predicted_class": predicted_class,
"confidence": confidence
})
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)
💬 Step 5: Test Your API
Run your FastAPI server:
uvicorn main:app --reload
Use tools like curl, Postman, or a simple HTML form to POST an image file to http://127.0.0.1:8000/predict/ and get back the predicted label with confidence score.
Example using curl:
curl -X POST -F "file=@/path/to/your/image.png" http://127.0.0.1:8000/predict/
📌 Notes:
✅ Conclusion
You've now built a complete AI-powered image classification API using TensorFlow and FastAPI. This setup is perfect for prototyping machine learning models and serving them quickly to applications or clients. With a bit more development, you can expand this into a scalable microservice or integrate with web/mobile interfaces for user-friendly AI-powered apps.