app.py

import io
import torch
from typing import Annotated
from PIL import Image
from fastapi import FastAPI, UploadFile, File
from torchvision import transforms
from torchvision.models import (
    mobilenet_v3_small,
    MobileNet_V3_Small_Weights
)
from db import init_db,save_prediction
import time

app = FastAPI(
    title="Edge Vision Inference Service"
)

weights = MobileNet_V3_Small_Weights.IMAGENET1K_V1
model = mobilenet_v3_small(weights=weights)
model.eval()

init_db()

categories = weights.meta["categories"]

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


def predict_image(file_bytes):
    image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
    input_tensor = preprocess(image).unsqueeze(0)

    with torch.no_grad():
        output = model(input_tensor)
        probabilities = torch.nn.functional.softmax(output[0], dim=0)

    top5_prob, top5_idx = torch.topk(probabilities, 5)

    predictions = []
    for prob, idx in zip(top5_prob, top5_idx):
        predictions.append({
            "class": categories[idx],
            "confidence": round(prob.item(), 4)
        })

    return predictions


@app.get("/")
def root():
    return {
        "message": "Service is running"
    }


@app.get("/status")
def status():
    return {
        "service": "online",
        "model": "loaded"
    }


@app.get("/model_info")
def model_info():
    return {
        "model": "MobileNetV3 Small",
        "weights": "ImageNet1K V1",
        "status": "loaded"
    }


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    start_time = time.time()
    file_bytes = await file.read()
    predictions = predict_image(file_bytes)
    latency_ms = (time.time() - start_time) * 1000
    top1 = predictions[0]
    save_prediction(
        filename=file.filename,
        top1_class=top1["class"],
        confidence=top1["confidence"],
        latency_ms=round(latency_ms, 2)
    )

    return {
        "filename": file.filename,
        "latency_ms": round(latency_ms, 2),
        "predictions": predictions
    }


@app.post("/predict_batch")
async def predict_batch(
    files: Annotated[list[UploadFile], File(description="Upload multiple images")]
):
    results = []
    for file in files:
        start_time = time.time()
        file_bytes = await file.read()
        predictions = predict_image(file_bytes)
        latency_ms = (time.time() - start_time) * 1000
        top1 = predictions[0]
        save_prediction(
            filename=file.filename,
            top1_class=top1["class"],
            confidence=top1["confidence"],
            latency_ms=round(latency_ms, 2)
        )
        results.append({
            "filename": file.filename,
            "latency_ms": round(latency_ms, 2),
            "predictions": predictions
        })

    return {
        "total_images": len(results),
        "results": results
    }

##uvicorn app:app --reload