""" Advanced FastAPI image-classification backend. Features -------- * Loads a modern, accurate CNN (MobileNet_V3_Large by default, EfficientNet_B0 optional) exactly once via an lru_cache-backed singleton. * Structured logging of every request, the predicted label and the wall-clock inference time. * /predict -> JSON: top-k classes, confidence scores, inference time (ms). * / -> Modern Tailwind single-page UI with drag-and-drop upload, live results and a loading spinner. * /health -> Lightweight readiness probe for cloud platforms. """ from __future__ import annotations import io import logging import os import time from functools import lru_cache from typing import List, Tuple import torch import torch.nn.functional as F from fastapi import FastAPI, File, HTTPException, UploadFile from fastapi.responses import HTMLResponse, JSONResponse from PIL import Image from torchvision import models, transforms # --------------------------------------------------------------------------- # # Configuration # --------------------------------------------------------------------------- # MODEL_NAME = os.getenv("MODEL_NAME", "mobilenet_v3_large").lower() # or efficientnet_b0 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") MAX_UPLOAD_BYTES = int(os.getenv("MAX_UPLOAD_BYTES", str(10 * 1024 * 1024))) # 10 MB # Keep the CPU + memory footprint small on free-tier hosts (e.g. Render's 512 MB). # Single-threaded inference uses noticeably less RAM and avoids CPU overcommit. torch.set_num_threads(int(os.getenv("TORCH_NUM_THREADS", "1"))) # --------------------------------------------------------------------------- # # Logging # --------------------------------------------------------------------------- # logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) logger = logging.getLogger("image-classifier") # --------------------------------------------------------------------------- # # Model + weights (cached so they load only once) # --------------------------------------------------------------------------- # @lru_cache(maxsize=1) def get_model_bundle() -> Tuple[torch.nn.Module, transforms.Compose, List[str]]: """Build, cache and return (model, preprocess, categories). The lru_cache guarantees the heavy weights are downloaded / loaded a single time for the lifetime of the process, regardless of how many requests hit the endpoint concurrently. """ logger.info("Loading model '%s' on device '%s' ...", MODEL_NAME, DEVICE) t0 = time.perf_counter() if MODEL_NAME == "efficientnet_b0": weights = models.EfficientNet_B0_Weights.IMAGENET1K_V1 model = models.efficientnet_b0(weights=weights) else: # Default: MobileNet_V3_Large — modern, fast and accurate. weights = models.MobileNet_V3_Large_Weights.IMAGENET1K_V2 model = models.mobilenet_v3_large(weights=weights) model.eval().to(DEVICE) preprocess = weights.transforms() # the canonical preprocessing for the weights categories = list(weights.meta["categories"]) logger.info( "Model ready in %.0f ms (%d classes).", (time.perf_counter() - t0) * 1000, len(categories), ) return model, preprocess, categories # --------------------------------------------------------------------------- # # Inference helpers # --------------------------------------------------------------------------- # def classify_bytes(raw: bytes, top_k: int = 5) -> Tuple[List[dict], float]: """Run inference on raw image bytes; return (results, inference_ms).""" model, preprocess, categories = get_model_bundle() try: image = Image.open(io.BytesIO(raw)).convert("RGB") except Exception as exc: # noqa: BLE001 raise HTTPException(status_code=400, detail=f"Invalid image file: {exc}") from exc tensor = preprocess(image).unsqueeze(0).to(DEVICE) t0 = time.perf_counter() with torch.inference_mode(): logits = model(tensor) probs = F.softmax(logits, dim=1)[0] inference_ms = (time.perf_counter() - t0) * 1000 top_k = max(1, min(top_k, probs.numel())) confidences, indices = torch.topk(probs, top_k) results = [ { "rank": rank, "label": categories[idx], "confidence": round(float(conf), 6), "confidence_pct": round(float(conf) * 100, 2), } for rank, (conf, idx) in enumerate(zip(confidences.tolist(), indices.tolist()), start=1) ] return results, inference_ms # --------------------------------------------------------------------------- # # FastAPI app # --------------------------------------------------------------------------- # app = FastAPI( title="Advanced Image Classification API", description="Production-ready PyTorch image classifier with a modern web UI.", version="1.0.0", ) @app.on_event("startup") def _warm_up() -> None: """Pre-load the model at boot so the first request isn't slow. On very small free-tier hosts you can set WARMUP=0 so the service becomes healthy immediately and the model loads lazily on the first request. """ if os.getenv("WARMUP", "1") == "1": get_model_bundle() @app.get("/health") def health() -> JSONResponse: return JSONResponse({"status": "ok", "model": MODEL_NAME, "device": str(DEVICE)}) @app.post("/predict") async def predict(file: UploadFile = File(...), top_k: int = 5) -> JSONResponse: """Classify an uploaded image and return the top-k predictions.""" request_start = time.perf_counter() if file.content_type and not file.content_type.startswith("image/"): raise HTTPException(status_code=400, detail="Uploaded file must be an image.") raw = await file.read() if not raw: raise HTTPException(status_code=400, detail="Empty file.") if len(raw) > MAX_UPLOAD_BYTES: raise HTTPException( status_code=413, detail=f"File too large (> {MAX_UPLOAD_BYTES // (1024 * 1024)} MB).", ) results, inference_ms = classify_bytes(raw, top_k=top_k) total_ms = (time.perf_counter() - request_start) * 1000 logger.info( "predict | file=%s size=%dB top1=%s (%.2f%%) inference=%.1fms total=%.1fms", file.filename, len(raw), results[0]["label"], results[0]["confidence_pct"], inference_ms, total_ms, ) return JSONResponse( { "filename": file.filename, "model": MODEL_NAME, "device": str(DEVICE), "inference_time_ms": round(inference_ms, 2), "total_time_ms": round(total_ms, 2), "predictions": results, } ) @app.get("/", response_class=HTMLResponse) def index() -> HTMLResponse: """Serve the single-page drag-and-drop UI.""" return HTMLResponse(INDEX_HTML) # --------------------------------------------------------------------------- # # Front-end (Tailwind via CDN, vanilla JS — zero build step) # --------------------------------------------------------------------------- # INDEX_HTML = """ VisionAI — Image Classifier

VisionAI

Image Classification

connecting…

Classify any image in milliseconds

Upload a photo and our deep-learning model will identify what it sees, ranked by confidence.

Drop your image here

or browse files

PNG · JPG · WEBP — up to 10 MB

Predictions

Your results will appear here once you upload an image.

Powered by PyTorch & FastAPI · runs locally on your machine

""" if __name__ == "__main__": import uvicorn uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True)