server.py

from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import JSONResponse

from nutrivision_api.model import build_model, get_device, load_image, predict

# Initialize the FastAPI application
app = FastAPI(
    title="NutriVision Food Classification API",
    description="Predict food class from an image using a trained ResNet50 model.",
    version="1.0.0"
)

# Load the model into the application state during startup to avoid reloading on every request
@app.on_event("startup")
def startup_event() -> None:
    device = get_device()
    app.state.device = device
    app.state.model = build_model(device)

# Health check endpoint to verify that the API is running and check the current hardware device
@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok", "device": str(app.state.device)}

# Prediction endpoint: accepts an image file and returns top 5 predictions
@app.post("/predict")
async def predict_endpoint(file: UploadFile = File(...)) -> JSONResponse:
    # Read the uploaded image file bytes
    content = await file.read()

    if not content:
        raise HTTPException(status_code=400, detail="Uploaded file is empty.")

    # Convert bytes to PIL Image object
    try:
        image = load_image(content)
    except Exception as exc:
        raise HTTPException(status_code=400, detail=f"Unable to read image file: {exc}")

    # Perform prediction using the loaded model
    predictions = predict(image, app.state.model, app.state.device, top_k=5)
    
    # Return the predictions as a JSON response
    return JSONResponse({
        "predictions": predictions,
        "top_prediction": predictions[0] if predictions else None
    })