Search⌘ K
AI Features
Home
Courses
Getting Started with Image Classification with PyTorch

REST API Integration

Explore how to deploy your trained PyTorch image classification model using FastAPI. Learn to create and run a REST API with HTTP POST requests for image inference, add CORS middleware, load models on startup, and test the server for prediction results.

We'll cover the following...

This lesson contains a step-by-step guide to integrate our image classification model into the FastAPI framework. We’ll serve it as an HTTP POST API.

We’ll reuse some of the code from our previous lesson to perform the following operations:

  • Create the FastAPI instance.
  • Add the CORS middleware.
Python
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
import torch
from PIL import Image
import torchvision.transforms as transforms
import timm
import io
app = FastAPI(
    title="Image Classification API",
    description="Working demo for weather classification.",
    version="0.0.1"
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
print('The title of the application is:', app.title)
print('The application contains the following middleware:', app.user_middleware)

Startup event

The startup event handler is ...