Spaces:

Sowaiba01
/

ThoraxNet

Sleeping

App Files Files Community

ThoraxNet / api /main.py

Sowaiba01

Update api/main.py

25127f2 verified 28 days ago

Raw

History Blame Contribute Delete

5.48 kB


	from __future__ import annotations
	import uuid
	from contextlib import asynccontextmanager
	from fastapi import FastAPI, Request
	from fastapi.middleware.cors import CORSMiddleware
	from fastapi.responses import HTMLResponse, JSONResponse
	from slowapi import Limiter, _rate_limit_exceeded_handler
	from slowapi.errors import RateLimitExceeded
	from slowapi.util import get_remote_address
	from api.inference import pipeline
	from api.routes import health, predict
	# ------------------------------------------------------------------
	# Rate limiter
	# ------------------------------------------------------------------
	limiter = Limiter(key_func=get_remote_address, default_limits=["10/minute"])
	# ------------------------------------------------------------------
	# Lifespan: startup / shutdown
	# ------------------------------------------------------------------
	@asynccontextmanager
	async def lifespan(app: FastAPI):
	print("[ChestAI] Loading model...")
	await pipeline.load()
	print("[ChestAI] Ready.")
	yield
	# Cleanup GradCAM hooks on shutdown.
	if pipeline.gradcam:
	pipeline.gradcam.remove_hooks()
	print("[ChestAI] Shutdown complete.")
	# ------------------------------------------------------------------
	# App
	# ------------------------------------------------------------------
	app = FastAPI(
	title="ChestAI",
	description=(
	"Uncertainty-aware multi-label chest X-ray diagnostic API. "
	"Analyzes 14 pathology classes with MC Dropout uncertainty estimation, "
	"GradCAM explainability, and auto-generated radiology reports."
	),
	version="1.0.0",
	lifespan=lifespan,
	docs_url="/docs",
	redoc_url="/redoc",
	)
	# Rate limiting
	app.state.limiter = limiter
	app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
	# CORS — allow frontend origins.
	# NOTE: allow_origins does NOT support wildcards like "https://*.vercel.app";
	# exact origins go in allow_origins, wildcard subdomains via allow_origin_regex.
	app.add_middleware(
	CORSMiddleware,
	allow_origins=[
	"http://localhost:3000",
	"http://localhost:3001",
	"https://thorax-tho.vercel.app", # production frontend
	],
	allow_origin_regex=r"https://.*\.vercel\.app", # Vercel preview deployments
	allow_credentials=True,
	allow_methods=["GET", "POST"],
	allow_headers=["*"],
	)
	# ------------------------------------------------------------------
	# Request ID middleware (for tracing in logs)
	# ------------------------------------------------------------------
	@app.middleware("http")
	async def add_request_id(request: Request, call_next):
	request_id = str(uuid.uuid4())[:8]
	request.state.request_id = request_id
	response = await call_next(request)
	response.headers["X-Request-ID"] = request_id
	return response
	# ------------------------------------------------------------------
	# Global exception handler
	# ------------------------------------------------------------------
	@app.exception_handler(Exception)
	async def global_exception_handler(request: Request, exc: Exception):
	return JSONResponse(
	status_code=500,
	content={"detail": "Internal server error.", "type": type(exc).__name__},
	)
	# ------------------------------------------------------------------
	# Landing page (registered BEFORE routers so it takes precedence over
	# any "/" route defined in a router)
	# ------------------------------------------------------------------
	LANDING_HTML = """<!DOCTYPE html>
	<html lang="en">
	<head>
	<meta charset="utf-8">
	<meta name="viewport" content="width=device-width, initial-scale=1">
	<title>ThoraxNet — AI Chest X-Ray Diagnosis</title>
	</head>
	<body style="margin:0;font-family:system-ui,-apple-system,sans-serif;background:#0b1220;color:#e8edf4;text-align:center;padding:8vh 1.5rem 3rem">
	<h1 style="font-size:2.4rem;margin-bottom:0.4rem">ThoraxNet</h1>
	<p style="color:#9fb0c3;max-width:36rem;margin:0 auto 2.2rem;line-height:1.6">
	AI chest X-ray diagnosis — detects <b>14 thoracic pathologies</b> with
	MC Dropout uncertainty estimates, GradCAM heatmaps and auto-generated
	radiology reports. Mean AUC <b>0.8215</b> on NIH ChestX-ray14.
	</p>
	<a href="https://thorax-tho.vercel.app" target="_blank" rel="noopener" style="display:inline-block;background:#22c55e;color:#04120a;font-weight:700;padding:0.9rem 2.4rem;border-radius:10px;text-decoration:none;font-size:1.15rem">
	Open Live Demo
	</a>
	<p style="margin-top:2.2rem">
	<a href="/docs" style="color:#7cc4ff;text-decoration:none">API Docs</a>
	·
	<a href="https://github.com/Sowaiba-01/ThoraxNet" target="_blank" rel="noopener" style="color:#7cc4ff;text-decoration:none">GitHub</a>
	·
	<a href="https://huggingface.co/Sowaiba01/chestai-model" target="_blank" rel="noopener" style="color:#7cc4ff;text-decoration:none">Model</a>
	</p>
	<p style="color:#5b6b7d;font-size:0.8rem;margin-top:3rem">
	For research use only. Not FDA cleared. Not a substitute for clinical radiologist interpretation.
	</p>
	</body>
	</html>"""


	@app.get("/", response_class=HTMLResponse, include_in_schema=False)
	async def root() -> HTMLResponse:
	return HTMLResponse(content=LANDING_HTML)
	# ------------------------------------------------------------------
	# Routers
	# ------------------------------------------------------------------
	app.include_router(health.router)
	app.include_router(predict.router)