Building a Production Vision AI App

From Script to Service

Everything we’ve built so far has been scripts: you run them manually, they process one image, print the result, and stop. That’s great for experimentation, but it’s not how software runs in production.

In production, your CV model needs to be a service, something that:

• Runs continuously, available 24/7
• Accepts requests from other systems (web apps, mobile apps, other APIs)
• Returns results in a structured format they can parse
• Handles errors gracefully without crashing
• Processes one request after another reliably

The standard way to build this in Python is a REST API: an HTTP server that accepts image uploads and returns JSON predictions. Other systems send a request, you return a response. Simple, universal, language-agnostic.

This module shows you how to build production-grade vision AI services.

Why Loading the Model Per-Request Is a Disaster

Before writing any code, let’s talk about the biggest performance mistake beginners make: loading the model inside the request handler.

# ❌ NEVER DO THIS
def predict_endpoint(image_file):
    model = load_model('my_model.pth')   # takes 5-10 seconds every time!
    return model.predict(image_file)

Loading a neural network from disk is expensive. The weights file for EfficientNet-B2 is 29MB. PyTorch needs to read it, allocate GPU memory, initialise layers, and warm up the compute graph. This takes 2-30 seconds depending on hardware.

If your API does this on every request, it will:

• Be unusably slow for users
• Exhaust server memory by loading the model repeatedly
• Time out on almost every request under any real load

The correct approach: load the model once when the server starts, then reuse it forever.

Building the API with FastAPI

FastAPI is Python’s most popular web framework for building APIs. It’s fast, has built-in data validation, and generates interactive documentation automatically.

pip install fastapi uvicorn python-multipart pillow torch torchvision

Start your application file:

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
import torch, torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import io, time

app = FastAPI(title="Vision AI API", version="1.0.0")

The app object is your application. FastAPI is a class from the fastapi library. The title and version appear in the auto-generated documentation at /docs.

Loading the Model at Startup

Use FastAPI’s startup event to load your model once when the server boots:

# Global variable: shared across all requests
model = None
CLASS_NAMES = ['defective', 'good', 'uncertain']

@app.on_event("startup")
async def load_model_on_startup():
    global model
    
    # Load the architecture
    m = models.efficientnet_b2(weights=None)
    
    # Replace the head to match your classes
    m.classifier[1] = nn.Linear(m.classifier[1].in_features, len(CLASS_NAMES))
    
    # Load your trained weights
    m.load_state_dict(torch.load('model.pth', map_location='cpu'))
    m.eval()
    
    model = m
    print("✓ Model loaded and ready")

Now define the preprocessing transform: this must match exactly what you used during training:

preprocess = transforms.Compose([
    transforms.Resize((260, 260)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

The Prediction Endpoint

Now define the endpoint that accepts image uploads:

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # Step 1: Validate the file type
    if not file.content_type or not file.content_type.startswith('image/'):
        raise HTTPException(
            status_code=400,
            detail=f"Expected an image file. Got: {file.content_type}"
        )
    
    # Step 2: Read the uploaded bytes
    contents = await file.read()

UploadFile is FastAPI’s type for uploaded files. File(...) means the field is required. The function is async so it doesn’t block while waiting for the upload.

    # Step 3: Validate file size (prevent abuse)
    if len(contents) > 10 * 1024 * 1024:   # 10MB
        raise HTTPException(status_code=400, detail="File too large. Maximum 10MB.")
    
    # Step 4: Decode the image
    try:
        img = Image.open(io.BytesIO(contents)).convert('RGB')
    except Exception:
        raise HTTPException(status_code=400, detail="Could not decode image.")

io.BytesIO turns the raw bytes into a file-like object that PIL.Image.open can read. .convert('RGB') ensures we always have 3 channels (handles grayscale and PNG with transparency).

    # Step 5: Run inference
    tensor = preprocess(img).unsqueeze(0)   # add batch dimension: (1, 3, 260, 260)
    
    start_time = time.time()
    with torch.no_grad():   # no gradients needed for inference
        outputs = model(tensor)
        probs   = torch.softmax(outputs, dim=1)[0]   # probabilities for each class
    latency_ms = (time.time() - start_time) * 1000
    
    # Step 6: Format and return the result
    predicted_idx  = probs.argmax().item()
    predicted_class = CLASS_NAMES[predicted_idx]
    
    return {
        "prediction":   predicted_class,
        "confidence":   round(probs[predicted_idx].item(), 4),
        "all_scores":   {name: round(probs[i].item(), 4) for i, name in enumerate(CLASS_NAMES)},
        "latency_ms":   round(latency_ms, 1),
        "model_version": "1.0.0",
    }

.unsqueeze(0) adds a batch dimension: the model expects input shape (batch, channels, height, width). Without it, you’d get shape errors.

Always return all_scores alongside the top prediction, the caller might want to apply their own threshold or understand the model’s uncertainty.

Running the Server

# Start the development server
uvicorn app:app --reload --host 0.0.0.0 --port 8000

--reload makes the server restart when you edit the code (development only). --host 0.0.0.0 makes it accessible from other machines on the network.

Visit http://localhost:8000/docs to get an interactive UI where you can upload test images directly in your browser.

Test with curl:

curl -X POST "http://localhost:8000/predict" \
     -F "file=@test_image.jpg"

Adding a Health Check

Production services need a health check endpoint: a simple route that confirms the server is running and the model is loaded:

@app.get("/health")
async def health_check():
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")
    return {"status": "healthy", "model_loaded": True}

Load balancers and monitoring systems call /health periodically. If it returns a non-200 response, they route traffic elsewhere and alert your team.

Containerising with Docker

Docker packages your application with all its dependencies into a portable container, it runs identically on your laptop, a cloud server, or a colleague’s machine.

Create a Dockerfile at the root of your project:

# Start from an official Python image
FROM python:3.11-slim

# Set the working directory inside the container
WORKDIR /app

# Copy and install dependencies first (Docker caches this layer)
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy the rest of your application
COPY . .

# Start the server when the container runs
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

Build and run:

docker build -t vision-api .
docker run -p 8000:8000 vision-api

-p 8000:8000 maps port 8000 in the container to port 8000 on your machine. After this, your API is accessible at http://localhost:8000/predict exactly as before, but running inside a self-contained container.

Exercise: Create a basic FastAPI server with a /predict endpoint. Load any pretrained model (even a simple ResNet-18 with imagenet weights). Test it with a few images from your browser via /docs. Add a /health endpoint. Measure the latency of predictions, is it consistent? Try submitting a non-image file and verify you get a proper 400 error.