""" MamaGuard -- FastAPI Server Maternal risk prediction API with clinical safety net and explainability. """ import torch import numpy as np import pickle import os from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse from api.schemas import PredictionRequest, PredictionResponse, AlertTier from api.alert_logic import compute_alert_tier, generate_action_text from src.model import MamaGuardMamba3 from src.explainability import explain_prediction from api.extract_report import router as extract_router # --- App setup ---------------------------------------------------------------- app = FastAPI( title="SheGuard -- Maternal Risk API", description=( "Predicts maternal mortality risk from prenatal visit sequences. " "Built with Mamba3 SSM for deployment in low-resource clinics." ), version="1.0.0", ) app.include_router(extract_router) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) # ─── Model loading ──────────────────────────────────────────────────────────── MODEL_PATH = "models/mamaguard_mamba3.pt" SCALER_PATH = "models/scaler.pkl" model = None scaler = None device = "cuda" if torch.cuda.is_available() else "cpu" FEATURE_ORDER = ['Age', 'SystolicBP', 'DiastolicBP', 'BS', 'BodyTemp', 'HeartRate'] FEATURE_DEFAULTS = { 'Age': 30.0, 'SystolicBP': 120.0, 'DiastolicBP': 80.0, 'BS': 7.5, 'BodyTemp': 36.8, 'HeartRate': 76.0 } @app.on_event("startup") async def load_model(): """Load model and scaler at server startup.""" global model, scaler if not os.path.exists(MODEL_PATH): print(f"WARNING: Model not found at {MODEL_PATH}. Run training first.") return model = MamaGuardMamba3(input_dim=6, d_model=64, n_layers=4, n_classes=3, d_state=32) model.load_state_dict(torch.load(MODEL_PATH, map_location=device)) model.to(device) model.eval() with open(SCALER_PATH, "rb") as f: scaler = pickle.load(f) print(f"MamaGuard model loaded on {device}") # --- Routes ------------------------------------------------------------------- @app.get("/health") async def health_check(): """Health check endpoint for load balancers.""" return { "status": "healthy", "model_loaded": model is not None, "device": device, "version": "1.0.0", } @app.post("/predict", response_model=PredictionResponse) async def predict(request: PredictionRequest): """ Process patient visit data and return risk prediction with explanation. """ if model is None or scaler is None: raise HTTPException( status_code=503, detail="Model not loaded. Contact system administrator." ) # Prepare visit data visits = request.visits visit_map = { 'Age': 'age', 'SystolicBP': 'systolic_bp', 'DiastolicBP': 'diastolic_bp', 'BS': 'blood_sugar', 'BodyTemp': 'body_temp', 'HeartRate': 'heart_rate' } raw_array = [] missing_counts = [] for visit in visits: row = [] missing = 0 for feat in FEATURE_ORDER: attr = visit_map[feat] val = getattr(visit, attr, None) if val is None: val = FEATURE_DEFAULTS[feat] missing += 1 row.append(val) raw_array.append(row) missing_counts.append(missing) raw_np = np.array(raw_array, dtype=np.float32) # Data quality score total_fields = len(visits) * len(FEATURE_ORDER) missing_total = sum(missing_counts) data_quality = round(1.0 - missing_total / total_fields, 3) # Scale scaled_np = scaler.transform(raw_np) # Pad or truncate to SEQ_LEN=5 SEQ_LEN = 5 n = len(scaled_np) if n < SEQ_LEN: pad = np.zeros((SEQ_LEN - n, scaled_np.shape[1]), dtype=np.float32) scaled_np = np.vstack([pad, scaled_np]) elif n > SEQ_LEN: scaled_np = scaled_np[-SEQ_LEN:] # Clinical safety net (hard WHO rules) from api.alert_logic import apply_clinical_safety_net, AlertTier visits_raw = [v.model_dump() for v in request.visits] forced_tier, forced_reason = apply_clinical_safety_net(visits_raw) # Run model and explain explanation = explain_prediction(model, scaled_np, scaler, device) if forced_tier is not None: explanation["top_reasons"].insert(0, f"[WHO guideline] {forced_reason}") # Alert tier alert_tier, suppressed = compute_alert_tier( probabilities = explanation["probabilities"], patient_id = request.patient_id, visit_importances= explanation["visit_importance"], ) # Override with clinical rule minimum if needed tier_order = {AlertTier.GREEN: 0, AlertTier.AMBER: 1, AlertTier.RED: 2} if forced_tier is not None: if tier_order[forced_tier] > tier_order[alert_tier]: alert_tier = forced_tier suppressed = False # Action text action, transfer_order = generate_action_text( tier = alert_tier, top_reasons = explanation["top_reasons"], data_quality = data_quality, staff_available = request.staff_available, blood_units = request.blood_units, ) return PredictionResponse( patient_id = request.patient_id, risk_level = explanation["risk_level"], alert_tier = alert_tier, confidence = explanation["confidence"], data_quality = data_quality, top_reasons = explanation["top_reasons"], feature_importance= explanation["feature_importance"], visit_importance = explanation["visit_importance"], action_required = action, transfer_order = transfer_order, suppressed = suppressed, probabilities = explanation["probabilities"], ) @app.get("/stats") async def get_stats(): """Returns patient assessment statistics.""" from api.alert_logic import _alert_history total = len(_alert_history) by_tier = {"GREEN": 0, "AMBER": 0, "RED": 0} for v in _alert_history.values(): by_tier[v["tier"].value] += 1 return { "patients_assessed": total, "by_tier": by_tier, "model_version": "SheGuard-Mamba3-v1", } # --- Dashboard static files --------------------------------------------------- DASHBOARD_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "dashboard") if os.path.isdir(DASHBOARD_DIR): app.mount("/static", StaticFiles(directory=DASHBOARD_DIR), name="dashboard-static") @app.get("/") async def serve_dashboard(): """Serve the dashboard HTML at root.""" index_path = os.path.join(DASHBOARD_DIR, "index.html") if os.path.exists(index_path): return FileResponse(index_path, media_type="text/html") return {"status": "ok", "model_loaded": model is not None, "device": device}