import io import os import tempfile import cv2 import numpy as np import traceback from fastapi import FastAPI, UploadFile, File, HTTPException, Query from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel from inference import run_inference, load_model import os app = FastAPI(title="DenseNet121-CBAM CT Scan API", version="1.0.0") # ─── Startup logging ─────────────────────────────────────────────────────── @app.on_event("startup") async def startup_event(): print("="*60) print(" CT Scan Classifier API - Ready to serve requests") print(" Model will load on first prediction request (lazy loading)") print(" Health endpoint: /health") print(" API docs: /docs") print("="*60) # ─── CORS — allow your React app origin ─────────────────────────────────────── ALLOWED_ORIGINS = os.getenv("ALLOWED_ORIGINS", "*").split(",") app.add_middleware( CORSMiddleware, allow_origins=ALLOWED_ORIGINS, # set to your Vercel URL in prod allow_credentials=True, allow_methods=["GET", "POST", "OPTIONS"], allow_headers=["*"], ) WEIGHTS_PATH = os.getenv("WEIGHTS_PATH", "trainedmodels/Model.pth") META_PATH = os.getenv("META_PATH", "trainedmodels/Model.json") DEVICE = os.getenv("DEVICE", "cpu") # ─── Lazy loading: model loads on first request, not at startup ───────────────────────────────────────────────── # This prevents timeout on HuggingFace free tier during container startup _model_loaded = False # ─── Routes ─────────────────────────────────────────────────────────────────── from fastapi.responses import HTMLResponse @app.get("/", response_class=HTMLResponse) def root(): return """

CT Scan API

Status: Running

📖 API Docs (Swagger)

❤️ Health Check

""" @app.get("/health") def health(): """Lightweight health check - does NOT load model (prevents timeout).""" return {"status": "ok", "device": DEVICE, "model": "DenseNet121-CBAM", "model_loaded": _model_loaded} @app.post("/predict") async def predict( file: UploadFile = File(...), gradcam: bool = Query(False), ): global _model_loaded # Ensure model is loaded on first request if not _model_loaded: try: load_model(WEIGHTS_PATH, DEVICE, meta_path=META_PATH) _model_loaded = True print("Model loaded and cached on first prediction request.") except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}") # Validate file type if not file.content_type.startswith("image/"): raise HTTPException(status_code=422, detail="Upload must be an image file.") # Max size guard — 10MB contents = await file.read() if len(contents) > 10 * 1024 * 1024: raise HTTPException(status_code=413, detail="Image must be under 10MB.") suffix = "." + file.filename.rsplit(".", 1)[-1].lower() with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(contents) tmp_path = tmp.name try: result = run_inference( tmp_path, WEIGHTS_PATH, device=DEVICE, generate_gradcam=gradcam ) except Exception as e: print("="*60) print("ERROR during inference:") print(traceback.format_exc()) print("="*60) raise HTTPException(status_code=500, detail=f"Inference failed: {str(e)}") finally: os.unlink(tmp_path) # If GradCAM requested, stream PNG back with prediction in headers if gradcam and result["gradcam_overlay"] is not None: overlay_bgr = cv2.cvtColor(result["gradcam_overlay"], cv2.COLOR_RGB2BGR) _, buf = cv2.imencode(".png", overlay_bgr) return StreamingResponse( io.BytesIO(buf.tobytes()), media_type="image/png", headers={ "X-Prediction": result["label"], "X-Probability": str(result["probability"]), "X-Threshold": str(result["threshold_used"]), "Access-Control-Expose-Headers": "X-Prediction,X-Probability,X-Threshold", }, ) result.pop("gradcam_overlay") return JSONResponse(result)