""" api/main.py ─────────── FastAPI layer exposing the breast cancer classifier as an HTTP API. Endpoints ───────── GET /health — Liveness check POST /predict — Single image inference (histopathology) POST /predict/batch — Multi-image batch inference POST /explain — Prediction + LLM natural language report POST /explain/visual — Prediction + Grad-CAM heatmap (base64 PNG) POST /chat — Conversational follow-up POST /mammogram/predict — Mammogram inference (3-model ensemble) POST /mammogram/explain — Mammogram + LLM explanation POST /mammogram/visual — Mammogram + Grad-CAM + LLM explanation Run locally ─────────── uvicorn api.main:app --host 0.0.0.0 --port 8000 --reload Env vars ──────── WEIGHTS_PATH model/weights.pth (histopathology) MAMMO_WEIGHTS_DIR Directory of ensemble members (default: model) MAMMO_THRESHOLD Ensemble decision threshold (default: 0.5) DEVICE "cuda" | "mps" | "cpu" (default: auto) CONFIDENCE_THR Float in [0,1] (default: 0.5) USE_LLM_MODEL "true" | "false" (default: false) LLM_MODEL_NAME HuggingFace model id (default: flan-t5-base) GRADCAM_ALPHA Heatmap blend opacity 0–1 (default: 0.5) """ from __future__ import annotations import base64 import io import logging import math import os import sys from pathlib import Path from typing import List, Optional import torch from fastapi import FastAPI, File, HTTPException, Query, UploadFile, status from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from PIL import Image # Ensure project root is on PYTHONPATH when running from repo root sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from model.inference import BreastCancerInferencePipeline from model.mammogram_inference import MammogramInferencePipeline from model.mammogram_ensemble import MammogramEnsemble from explainability.gradcam import GradCAM from explainability.mammogram_gradcam import MammogramGradCAM from explainability.chat_pipeline import DualModelChatPipeline, create_pipeline from explainability.llm_explain import LLMExplainer, ChatEngine from explainability.llm_chat import stream_chat, llm_available # ── Logging ─────────────────────────────────────────────────────────────────── logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # ── FastAPI app ─────────────────────────────────────────────────────────────── app = FastAPI( title="MedAI — Breast Cancer Analysis Platform", description=( "DenseNet-121 histopathology classifier and 3-model EfficientNet-B4 " "mammogram ensemble, with Grad-CAM explainability and LLM explanation. " "Research and educational use only. Not a standalone diagnostic tool." ), version="2.1.0", docs_url="/docs", redoc_url="/redoc", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["GET", "POST"], allow_headers=["*"], ) # ── Module singletons ───────────────────────────────────────────────────────── _pipeline: BreastCancerInferencePipeline | None = None _mammo_pipeline: "MammogramEnsemble | MammogramInferencePipeline | None" = None _gradcam: GradCAM | None = None _llm_explainer: LLMExplainer | None = None _chat_engine: ChatEngine | None = None _chat_pipeline: DualModelChatPipeline | None = None @app.on_event("startup") def _load_pipeline() -> None: global _pipeline, _mammo_pipeline weights = os.getenv("WEIGHTS_PATH", "model/weights.pth") device = os.getenv("DEVICE", None) thr = float(os.getenv("CONFIDENCE_THR", "0.5")) weights_path = Path(weights) if Path(weights).exists() else None if weights_path is None: logger.warning( "weights.pth not found at '%s'. " "Running with ImageNet-pretrained backbone only.", weights, ) _pipeline = BreastCancerInferencePipeline( weights_path=weights_path, device=device, confidence_threshold=thr, ) logger.info("Histopathology pipeline ready on device: %s", _pipeline.device) # ── Mammogram: prefer the 3-model ensemble; fall back to single-view ────── mammo_dir = Path(os.getenv("MAMMO_WEIGHTS_DIR", "model")) mammo_thr = float(os.getenv("MAMMO_THRESHOLD", "0.5")) members = [mammo_dir / f"model_s{s}.pth" for s in (42, 123, 999)] present = [p for p in members if p.exists()] if present: _mammo_pipeline = MammogramEnsemble( weight_paths = present, device = device, threshold = mammo_thr, ) logger.info( "Mammogram ENSEMBLE ready on %s (%d members, threshold=%.2f)", _mammo_pipeline.device, len(present), mammo_thr, ) else: legacy = mammo_dir / "mammogram_weights.pth" _mammo_pipeline = MammogramInferencePipeline( weights_path = legacy if legacy.exists() else None, device = device, ) logger.warning( "Ensemble members not found in '%s'. Falling back to single-view " "pipeline (%s).", mammo_dir, "trained weights" if legacy.exists() else "ImageNet init", ) logger.info("Grad-CAM and LLM explainer load on first /explain request.") def _get_mammo_pipeline(): if _mammo_pipeline is None: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Mammogram pipeline not ready.", ) return _mammo_pipeline def _get_pipeline() -> BreastCancerInferencePipeline: if _pipeline is None: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Model pipeline is not ready. Try again shortly.", ) return _pipeline def _get_gradcam() -> GradCAM: """Lazy-load GradCAM on first explainability request (histopathology).""" global _gradcam if _gradcam is None: pipeline = _get_pipeline() alpha = float(os.getenv("GRADCAM_ALPHA", "0.5")) logger.info("Initialising Grad-CAM module (device=%s, alpha=%.2f)...", pipeline.device, alpha) _gradcam = GradCAM( model = pipeline.model, device = str(pipeline.device), alpha = alpha, ) logger.info("Grad-CAM ready.") return _gradcam def _get_llm() -> LLMExplainer: global _llm_explainer if _llm_explainer is None: use_llm = os.getenv("USE_LLM_MODEL", "false").lower() == "true" model_name = os.getenv("LLM_MODEL_NAME", "google/flan-t5-base") logger.info( "Initialising LLM explainer (use_llm=%s, model=%s)...", use_llm, model_name if use_llm else "template-engine", ) _llm_explainer = LLMExplainer(model_name=model_name, use_llm=use_llm) logger.info("LLM explainer ready (engine=%s).", "flan-t5" if use_llm else "template") return _llm_explainer def _get_chat_engine() -> ChatEngine: global _chat_engine if _chat_engine is None: _chat_engine = ChatEngine() logger.info("ChatEngine fallback ready.") return _chat_engine def _get_chat_pipeline() -> "DualModelChatPipeline | None": global _chat_pipeline if _chat_pipeline is None: _chat_pipeline = create_pipeline() if _chat_pipeline: logger.info("Dual pipeline ready: BioMedLM + Llama 3.2 (Groq)") else: logger.info( "Dual pipeline not available — set GROQ_API_KEY + HF_TOKEN " "to enable BioMedLM + Llama 3.2. Using ChatEngine fallback." ) return _chat_pipeline # ── Response schemas ────────────────────────────────────────────────────────── class PredictionResponse(BaseModel): prediction: str = Field(..., examples=["malignant"]) confidence: float = Field(..., ge=0.0, le=1.0, examples=[0.875]) logits: list[float]= Field( ..., description="Raw pre-softmax scores [benign, malignant]", examples=[[-2.14, 3.87]], ) disclaimer: str = Field( default=("Research / educational use only. " "Not a standalone diagnostic tool.") ) class ExplainResponse(PredictionResponse): summary: str = Field(..., description="Plain-language summary of the prediction") detail: str = Field(..., description="Deeper explanation with confidence context") audience: str = Field(..., description="Target audience the explanation is tailored for") engine: str = Field(..., description="'flan-t5' if LLM ran, 'template' if fallback") class VisualExplainResponse(ExplainResponse): overlay_b64: str = Field(..., description="Base64 PNG of Grad-CAM overlay") heatmap_b64: str = Field(..., description="Base64 PNG of raw Grad-CAM heatmap") spatial_summary: str = Field(..., description="Regions that drove the prediction") heatmap_mean: float = Field(..., description="Mean activation value [0,1]") heatmap_max: float = Field(..., description="Peak activation value [0,1]") class HealthResponse(BaseModel): status: str device: str gradcam_loaded: bool llm_loaded: bool llm_engine: str mammogram_mode: str chat_llm: str # ── Helpers ─────────────────────────────────────────────────────────────────── ACCEPTED_MIME = {"image/jpeg", "image/png", "image/tiff", "image/bmp"} def _validate_and_open(upload: UploadFile) -> Image.Image: if upload.content_type not in ACCEPTED_MIME: raise HTTPException( status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, detail=f"Unsupported file type: {upload.content_type}. " f"Accepted: {', '.join(ACCEPTED_MIME)}", ) data = upload.file.read() try: return Image.open(io.BytesIO(data)).convert("RGB") except Exception as exc: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=f"Could not decode image: {exc}", ) from exc def _serialize(result: dict) -> dict: return { "prediction": result["prediction"], "confidence": result["confidence"], "logits": result["logits"].squeeze().tolist(), } def _pil_to_b64(img: Image.Image, fmt: str = "PNG") -> str: buf = io.BytesIO() img.save(buf, format=fmt) return base64.b64encode(buf.getvalue()).decode("utf-8") def _heatmap_to_pil(heatmap) -> Image.Image: import numpy as np arr = (heatmap * 255).astype(np.uint8) return Image.fromarray(arr, mode="L") def _ensemble_logits(result: dict) -> list[float]: """ The ensemble returns averaged probabilities, not raw logits. Synthesize log-probabilities so the response schema and the /chat context (which expects [benign, malignant] scores) stay backward-compatible. """ mal = float(result.get("malignant_probability", result.get("confidence", 0.5))) ben = 1.0 - mal return [round(math.log(max(ben, 1e-6)), 6), round(math.log(max(mal, 1e-6)), 6)] def _is_ensemble(pipeline) -> bool: return isinstance(pipeline, MammogramEnsemble) class PatientRecord(BaseModel): name: str = Field(default="", description="Patient name") age: int = Field(default=0, description="Patient age") sex: str = Field(default="", description="Male / Female / Other") medical_history: str = Field(default="", description="Relevant medical history") symptoms: str = Field(default="", description="Current symptoms or concerns") previous_scans: str = Field(default="", description="Previous imaging history") class ChatRequest(BaseModel): message: str audience: str = Field(default="clinician", description="clinician | researcher | patient") prediction: str = Field(default="") confidence: float = Field(default=0.0, ge=0.0, le=1.0) logits: list[float] = Field(default=[0.0, 0.0]) spatial_summary: str = Field(default="") history: list = Field(default_factory=list) patient: PatientRecord = Field(default_factory=PatientRecord) class ChatResponse(BaseModel): response: str audience: str # ── Endpoints ───────────────────────────────────────────────────────────────── @app.get("/health", response_model=HealthResponse, tags=["Meta"]) def health() -> HealthResponse: pipeline = _get_pipeline() llm_engine = "not_loaded" if _llm_explainer is not None: llm_engine = "flan-t5" if _llm_explainer.use_llm else "template" mammo_mode = "not_loaded" if _mammo_pipeline is not None: mammo_mode = "ensemble" if _is_ensemble(_mammo_pipeline) else "single-view" import os avail, provider = llm_available() model = os.getenv("CHAT_MODEL") or ("openai/gpt-oss-120b" if provider == "groq" else provider) chat_llm = f"{provider}/{model}" if avail else f"{provider}/no_api_key" return HealthResponse( status = "ok", device = str(pipeline.device), gradcam_loaded = _gradcam is not None, llm_loaded = _llm_explainer is not None, llm_engine = llm_engine, mammogram_mode = mammo_mode, chat_llm = chat_llm, ) @app.post("/predict", response_model=PredictionResponse, tags=["Inference"]) def predict(file: UploadFile = File(...)) -> PredictionResponse: pipeline = _get_pipeline() image = _validate_and_open(file) try: result = pipeline.predict(image) except Exception as exc: logger.exception("Inference error") raise HTTPException(500, detail=f"Inference failed: {exc}") from exc return PredictionResponse(**_serialize(result)) @app.post("/predict/batch", response_model=List[PredictionResponse], tags=["Inference"]) def predict_batch(files: List[UploadFile] = File(...)) -> list[PredictionResponse]: if len(files) > 16: raise HTTPException(400, detail="Batch size exceeds maximum of 16 images.") pipeline = _get_pipeline() images = [_validate_and_open(f) for f in files] try: results = pipeline.predict_batch(images) except Exception as exc: logger.exception("Batch inference error") raise HTTPException(500, detail=f"Batch inference failed: {exc}") from exc return [PredictionResponse(**_serialize(r)) for r in results] @app.post("/explain", response_model=ExplainResponse, tags=["Explainability"]) def explain( file: UploadFile = File(...), audience: str = Query(default="clinician", enum=["clinician", "researcher", "patient"]), ) -> ExplainResponse: pipeline = _get_pipeline() llm = _get_llm() image = _validate_and_open(file) try: prediction = pipeline.predict(image) except Exception as exc: logger.exception("Inference error in /explain") raise HTTPException(500, detail=f"Inference failed: {exc}") from exc try: report = llm.explain(prediction, audience=audience) except Exception as exc: logger.exception("LLM explanation error in /explain") raise HTTPException(500, detail=f"Explanation generation failed: {exc}") from exc return ExplainResponse( **_serialize(prediction), summary = report["summary"], detail = report["detail"], audience = report["audience"], engine = report["engine"], ) @app.post("/explain/visual", response_model=VisualExplainResponse, tags=["Explainability"]) def explain_visual( file: UploadFile = File(...), audience: str = Query(default="clinician", enum=["clinician", "researcher", "patient"]), class_idx: Optional[int] = Query(default=None, ge=0, le=1), ) -> VisualExplainResponse: llm = _get_llm() cam = _get_gradcam() image = _validate_and_open(file) try: cam_result = cam.explain(image, class_idx=class_idx) except Exception as exc: logger.exception("Grad-CAM error in /explain/visual") raise HTTPException(500, detail=f"Grad-CAM failed: {exc}") from exc try: report = llm.explain_with_gradcam(cam_result, audience=audience) except Exception as exc: logger.exception("LLM explanation error in /explain/visual") raise HTTPException(500, detail=f"Explanation generation failed: {exc}") from exc import numpy as np heatmap = cam_result["heatmap"] overlay = cam_result["overlay"] return VisualExplainResponse( **_serialize(cam_result), summary = report["summary"], detail = report["detail"], audience = report["audience"], engine = report["engine"], overlay_b64 = _pil_to_b64(overlay), heatmap_b64 = _pil_to_b64(_heatmap_to_pil(heatmap)), spatial_summary = LLMExplainer._summarise_heatmap(heatmap), heatmap_mean = round(float(np.mean(heatmap)), 6), heatmap_max = round(float(np.max(heatmap)), 6), ) @app.post("/chat", response_model=ChatResponse, tags=["Explainability"]) def chat(request: ChatRequest) -> ChatResponse: logits = request.logits if len(request.logits) >= 2 else [0.0, 0.0] patient_dict = {} if request.patient: patient_dict = { "name": request.patient.name, "age": request.patient.age, "sex": request.patient.sex, "medical_history": request.patient.medical_history, "symptoms": request.patient.symptoms, "previous_scans": request.patient.previous_scans, } pipeline = _get_chat_pipeline() if pipeline: try: response_text = pipeline.respond( message = request.message, audience = request.audience, prediction = request.prediction, confidence = request.confidence, benign_logit = logits[0], malignant_logit = logits[1], spatial_summary = request.spatial_summary, history = request.history, patient = patient_dict, ) if response_text: return ChatResponse(response=response_text, audience=request.audience) except Exception as e: logger.warning("Dual pipeline error (%s) — falling back to ChatEngine.", e) _get_llm() engine = _get_chat_engine() response_text = engine.respond( message = request.message, audience = request.audience, prediction = request.prediction, confidence = request.confidence, benign_logit = logits[0], malignant_logit = logits[1], spatial_summary = request.spatial_summary, history = request.history, patient = patient_dict, ) return ChatResponse(response=response_text, audience=request.audience) class LLMChatRequest(BaseModel): messages: list = Field(default_factory=list, description="[{role: 'user'|'assistant', content: str}, ...]") audience: str = Field(default="clinician") prediction: str = Field(default="") confidence: float = Field(default=0.0, ge=0.0, le=1.0) logits: list[float] = Field(default=[0.0, 0.0]) spatial_summary: str = Field(default="") birads: str = Field(default="") @app.post("/chat/stream", tags=["Explainability"]) def chat_stream(req: LLMChatRequest): """Real LLM assistant with token streaming (Claude or GPT via env config).""" context = { "prediction": req.prediction, "confidence": req.confidence, "logits": req.logits, "spatial_summary": req.spatial_summary, "birads": req.birads, } if req.prediction else None def generate(): try: for token in stream_chat(req.messages, context=context, audience=req.audience): yield token except Exception as e: # noqa: BLE001 logger.warning("LLM chat stream error: %s", e) yield f"\n[Assistant error: {e}]" return StreamingResponse(generate(), media_type="text/plain; charset=utf-8") class MammogramResponse(BaseModel): prediction: str confidence: float logits: list[float] birads: str modality: str = "mammogram" per_model: Optional[list[float]] = Field( default=None, description="Per-member malignant probabilities (ensemble mode only)", ) n_models: Optional[int] = Field( default=None, description="Number of ensemble members averaged", ) disclaimer: str = ( "AI-assisted mammogram analysis. Research use only. " "Not a substitute for radiologist review or clinical diagnosis." ) class MammogramExplainResponse(MammogramResponse): summary: str detail: str audience: str engine: str class MammogramVisualResponse(MammogramExplainResponse): overlay_b64: str heatmap_b64: str spatial_summary: str heatmap_mean: float heatmap_max: float def _run_mammo_predict(pipeline, image: Image.Image, tta: bool) -> dict: """Unified prediction → dict with prediction/confidence/logits/birads/etc.""" if _is_ensemble(pipeline): result = pipeline.predict_tta(image) if tta else pipeline.predict(image) return { "prediction": result["prediction"], "confidence": result["confidence"], "logits": _ensemble_logits(result), "birads": result["birads"], "modality": result.get("modality", "mammogram_ensemble"), "per_model": result.get("per_model"), "n_models": result.get("n_models"), } # Legacy single-view pipeline (returns logits tensor) result = pipeline.predict(image) return { "prediction": result["prediction"], "confidence": result["confidence"], "logits": result["logits"].squeeze().tolist(), "birads": result["birads"], "modality": result.get("modality", "mammogram"), "per_model": None, "n_models": None, } @app.post("/mammogram/predict", response_model=MammogramResponse, tags=["Mammogram"]) def mammogram_predict( file: UploadFile = File(...), tta: bool = Query(default=False, description="Test-time augmentation (ensemble only) — " "slower, slightly more accurate"), ) -> MammogramResponse: """ Classify a mammogram using the 3-model EfficientNet-B4 ensemble. Returns prediction, confidence, BI-RADS suggestion, and per-member probs. """ pipeline = _get_mammo_pipeline() image = _validate_and_open(file) try: result = _run_mammo_predict(pipeline, image, tta) except Exception as exc: logger.exception("Mammogram inference error") raise HTTPException(500, detail=f"Mammogram inference failed: {exc}") from exc return MammogramResponse(**result) @app.post("/mammogram/explain", response_model=MammogramExplainResponse, tags=["Mammogram"]) def mammogram_explain( file: UploadFile = File(...), audience: str = Query(default="clinician", enum=["clinician", "researcher", "patient"]), tta: bool = Query(default=False), ) -> MammogramExplainResponse: """Classify a mammogram (ensemble) and generate an audience-specific explanation.""" pipeline = _get_mammo_pipeline() llm = _get_llm() image = _validate_and_open(file) try: result = _run_mammo_predict(pipeline, image, tta) except Exception as exc: raise HTTPException(500, detail=f"Mammogram inference failed: {exc}") from exc explain_input = { "prediction": result["prediction"], "confidence": result["confidence"], "logits": torch.tensor([result["logits"]]), "birads": result["birads"], } try: report = llm.explain(explain_input, audience=audience, modality="mammogram") except Exception as exc: raise HTTPException(500, detail=f"Explanation failed: {exc}") from exc return MammogramExplainResponse( **result, summary = report["summary"], detail = report["detail"], audience = report["audience"], engine = report["engine"], ) @app.post("/mammogram/visual", response_model=MammogramVisualResponse, tags=["Mammogram"]) def mammogram_visual( file: UploadFile = File(...), audience: str = Query(default="clinician", enum=["clinician", "researcher", "patient"]), class_idx: Optional[int] = Query(default=None, ge=0, le=1), ) -> MammogramVisualResponse: """ Full mammogram analysis: ensemble prediction + Grad-CAM + LLM explanation. Grad-CAM is inherently a single-model technique, so the heatmap is computed on the strongest ensemble member (the most representative saliency map), while the prediction label/confidence uses the full ensemble average. """ import numpy as np pipeline = _get_mammo_pipeline() llm = _get_llm() image = _validate_and_open(file) # Authoritative label/confidence from the full ensemble try: pred = _run_mammo_predict(pipeline, image, tta=False) except Exception as exc: raise HTTPException(500, detail=f"Mammogram inference failed: {exc}") from exc # Grad-CAM on the representative single model (EfficientNet-specific) cam_model = pipeline.cam_model if _is_ensemble(pipeline) else pipeline.model try: alpha = float(os.getenv("GRADCAM_ALPHA", "0.5")) mammo_cam = MammogramGradCAM(model=cam_model, device=str(pipeline.device), alpha=alpha) cam_result = mammo_cam.explain(image, class_idx=class_idx) except Exception as exc: logger.exception("Mammogram Grad-CAM error") raise HTTPException( 500, detail=(f"Grad-CAM failed: {exc}. The /mammogram/predict endpoint " f"still works without the heatmap."), ) from exc # Make the LLM explanation reflect the ENSEMBLE verdict (authoritative), # while the heatmap stays the representative member's saliency map. cam_result["prediction"] = pred["prediction"] cam_result["confidence"] = pred["confidence"] cam_result["logits"] = torch.tensor([pred["logits"]]) cam_result["birads"] = pred["birads"] try: report = llm.explain_with_gradcam(cam_result, audience=audience, modality="mammogram") except Exception as exc: raise HTTPException(500, detail=f"Explanation failed: {exc}") from exc heatmap = cam_result["heatmap"] overlay = cam_result["overlay"] return MammogramVisualResponse( prediction = pred["prediction"], confidence = pred["confidence"], logits = pred["logits"], birads = pred["birads"], modality = pred["modality"], per_model = pred["per_model"], n_models = pred["n_models"], summary = report["summary"], detail = report["detail"], audience = report["audience"], engine = report["engine"], overlay_b64 = _pil_to_b64(overlay), heatmap_b64 = _pil_to_b64(_heatmap_to_pil(heatmap)), spatial_summary = LLMExplainer._summarise_heatmap(heatmap), heatmap_mean = round(float(np.mean(heatmap)), 6), heatmap_max = round(float(np.max(heatmap)), 6), )