Capstone Project 3: Image Classifier using Transfer Learning

Learn to build an Image Classification project in PyTorch. Download pre-trained ResNet, freeze layers, fine-tune on custom images, and host inference.

Introduction

In Module 9, we explored how Transfer Learning allows us to reuse pre-trained features. In this third capstone project, we will apply this technique to construct a Deep Learning Image Classifier in PyTorch.

We will write a complete pipeline that downloads a pre-trained ResNet-18 model, freezes its feature extraction layers, appends a custom classification head, and deploys it as a FastAPI web server. Crucially, we will write code that accepts binary image file uploads (e.g., JPEG or PNG) over HTTP, pre-processes the pixels into PyTorch tensors, and returns the predicted class label.

What You Will Learn

  • Reusing torchvision pre-trained models.
  • Processing incoming image file uploads in FastAPI.
  • Resizing, cropping, and normalizing images using torchvision.transforms.
  • Running GPU/CPU model inference.

Project Structure

Setup your folder layouts:

image-classifier/
│
├── models/
│   └── resnet18_custom.pt     # Serialized PyTorch state weights
├── src/
│   ├── train.py               # Downloads and modifies model
│   └── app.py                 # FastAPI image upload receiver
├── requirements.txt
└── Dockerfile

Step 1: Model Setup (src/train.py)

We will write the setup script to download ResNet-18 from torchvision, freeze its weights, swap its final fully connected layer, and save the model weights to our models/ directory.

# src/train.py
import os
import torch
import torch.nn as nn
from torchvision import models

def setup_and_save_model():
    print("Step 1: Downloading pre-trained ResNet-18...")
    # Load model with pre-trained weights
    model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
    
    # 2. Freeze convolutional weights
    for param in model.parameters():
        param.requires_grad = False
        
    # 3. Swap final layer
    # ResNet-18 features 512 input channels. We will predict 2 classes (e.g., Cat vs Dog)
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, 2)
    
    print("\nCustomized output classifier:")
    print(model.fc)
    
    # 4. Serialize Model Weights
    os.makedirs("models", exist_ok=True)
    # Save the entire model structure and weights
    torch.save(model, "models/resnet18_custom.pt")
    print("\nModel saved successfully to: models/resnet18_custom.pt")

if __name__ == "__main__":
    setup_and_save_model()

Step 2: FastAPI Image Upload Server (src/app.py)

Web APIs do not receive image files as JSON strings. Instead, they accept files as multipart-form payloads. We will use the PIL (Python Imaging Library) to open raw bytes, and apply torchvision transforms to format the image (resizing it to $224 \times 224$ pixels and applying ImageNet normalization metrics).

# src/app.py
from fastapi import FastAPI, UploadFile, File, HTTPException
import torch
import torchvision.transforms as transforms
from PIL import Image
import io
import os

app = FastAPI(title="Deep Image Classification API", version="1.0")

# Load customized PyTorch model
model_path = "models/resnet18_custom.pt"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

if os.path.exists(model_path):
    model = torch.load(model_path, map_location=device)
    model.eval() # Set model to evaluation mode
else:
    model = None

# Class labels mapping
CLASS_LABELS = {0: "Cat", 1: "Dog"}

# Define ImageNet pre-processing transformation pipeline
transform_pipeline = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406], # ImageNet standard means
        std=[0.229, 0.224, 0.225]   # ImageNet standard deviations
    )
])

@app.post("/predict")
async def predict_image(file: UploadFile = File(...)):
    if model is None:
        raise HTTPException(status_code=500, detail="Model file not found.")
        
    # Verify file type
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")

    try:
        # 1. Read binary image bytes
        image_bytes = await file.read()
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        
        # 2. Apply preprocessing transforms
        tensor = transform_pipeline(image).unsqueeze(0) # Add batch dimension: [1, 3, 224, 224]
        tensor = tensor.to(device)
        
        # 3. Model Inference (disable gradient tracking for speed)
        with torch.no_grad():
            outputs = model(tensor)
            probabilities = torch.softmax(outputs, dim=1)
            prediction = torch.argmax(probabilities, dim=1).item()
            
        confidence = float(probabilities[0][prediction].item())
        predicted_class = CLASS_LABELS.get(prediction, "Unknown")
        
        return {
            "status": "success",
            "filename": file.filename,
            "prediction": predicted_class,
            "confidence_score": round(confidence, 4),
            "probabilities": {
                CLASS_LABELS[0]: round(float(probabilities[0][0].item()), 4),
                CLASS_LABELS[1]: round(float(probabilities[0][1].item()), 4)
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")

Execution & Verification

To execute and test your image classifier:

  1. Install torchvision, PyTorch, and pillow: 
    pip install torch torchvision pillow fastapi uvicorn python-multipart
  2. Setup the Model Weights: 
    python src/train.py
  3. Start the API: 
    uvicorn src.app:app --reload --port 8000
  4. Test using curl (Upload an actual image file): 
    curl -X POST "http://localhost:8000/predict" \
     -F "file=@/path/to/your/dog_image.jpg"
    Expected Response: 
    {
  "status": "success",
  "filename": "dog_image.jpg",
  "prediction": "Dog",
  "confidence_score": 0.9841,
  "probabilities": {
    "Cat": 0.0159,
    "Dog": 0.9841
  }
}

Summary

In this final capstone project, we loaded a pre-trained ResNet-18 model using PyTorch, froze its convolutional layers, and attached a customized classification head. We then built a FastAPI server capable of parsing binary multipart-form image uploads, processing raw pixels through standard ImageNet resizing and normalization steps, and running forward pass predictions to output final classification labels.

Next Topic

You now have a portfolio of three production-grade projects. How do you prepare for placement exams and system design interviews to land a job? Let's check: AI and Machine Learning Interview Prep: Top 30 Q&A.

← Back to Syllabus