Spaces:
Sleeping
Sleeping
| 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 ─────────────────────────────────────────────────────── | |
| 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 | |
| def root(): | |
| return """ | |
| <html> | |
| <body style="font-family:sans-serif; text-align:center; padding:40px"> | |
| <h2> CT Scan API</h2> | |
| <p>Status: <strong style="color:green">Running</strong></p> | |
| <p><a href="/docs">📖 API Docs (Swagger)</a></p> | |
| <p><a href="/health">❤️ Health Check</a></p> | |
| </body> | |
| </html> | |
| """ | |
| def health(): | |
| """Lightweight health check - does NOT load model (prevents timeout).""" | |
| return {"status": "ok", "device": DEVICE, "model": "DenseNet121-CBAM", "model_loaded": _model_loaded} | |
| 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) | |