Spaces:
Sleeping
Sleeping
File size: 4,999 Bytes
78c0a02 2910e7f 78c0a02 12b6b94 78c0a02 43a582e 78c0a02 12b6b94 78c0a02 f973701 78c0a02 f0b4428 78c0a02 43a582e 78c0a02 f973701 78c0a02 2910e7f 78c0a02 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | 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 """
<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>
"""
@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)
|