Spaces:
Runtime error
Runtime error
| """ | |
| main.py — FastAPI backend for Radiology AI Report Generation | |
| Vision is ENABLED by default (BLIP-2 reads the actual image). | |
| """ | |
| import io | |
| import json | |
| import logging | |
| import os | |
| import sys | |
| import uuid | |
| from datetime import datetime | |
| from pathlib import Path | |
| from typing import Optional | |
| ROOT = Path(__file__).resolve().parent.parent | |
| if str(ROOT) not in sys.path: | |
| sys.path.insert(0, str(ROOT)) | |
| from fastapi import FastAPI, File, Form, HTTPException, UploadFile | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse, Response | |
| # Database | |
| import importlib.util as _ilu | |
| _db_spec = _ilu.spec_from_file_location("backend.database", | |
| __file__.replace("main.py","database.py")) | |
| _db_mod = _ilu.module_from_spec(_db_spec) | |
| _db_spec.loader.exec_module(_db_mod) | |
| db = _db_mod | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s | %(levelname)-8s | %(name)s — %(message)s", | |
| ) | |
| logger = logging.getLogger(__name__) | |
| app = FastAPI( | |
| title="Radiology AI Report Generation System", | |
| version="2.0.0", | |
| description="AI-powered radiology report generation with BLIP-2 vision + Phi-2 LLM", | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| _report_store: dict = {} | |
| _models_loaded = False | |
| _slm_model = None | |
| _slm_tokenizer = None | |
| _vision_model = None | |
| _vision_processor = None | |
| # Vision is ON by default — set ENABLE_VISION_CAPTION=false to disable | |
| ENABLE_VISION_CAPTION = False | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| # Module loader helper (avoids import path issues on Windows) | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| def _load_module(name, filepath): | |
| import importlib.util | |
| if name in sys.modules: | |
| return sys.modules[name] | |
| spec = importlib.util.spec_from_file_location(name, filepath) | |
| mod = importlib.util.module_from_spec(spec) | |
| sys.modules[name] = mod | |
| spec.loader.exec_module(mod) | |
| return mod | |
| _BACKEND = Path(__file__).resolve().parent | |
| def _ml(): | |
| return _load_module("backend.model_loader", _BACKEND / "model_loader.py") | |
| def _pp(): | |
| return _load_module("backend.preprocessing", _BACKEND / "preprocessing.py") | |
| def _rg(): | |
| return _load_module("backend.report_generator",_BACKEND / "report_generator.py") | |
| def _vm(): | |
| return _load_module("backend.validation_module",_BACKEND / "validation_module.py") | |
| def _mc(): | |
| return _load_module("backend.medical_coding", _BACKEND / "medical_coding.py") | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| # Model loading | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| def _ensure_models_loaded(): | |
| global _models_loaded, _slm_model, _slm_tokenizer, _vision_model, _vision_processor | |
| if _models_loaded: | |
| return | |
| logger.info("Loading AI models...") | |
| ml = _ml() | |
| # Always load language model | |
| _slm_model, _slm_tokenizer = ml.load_slm() | |
| # Load vision model (BLIP-2) if enabled | |
| if ENABLE_VISION_CAPTION: | |
| logger.info("Loading BLIP-2 vision model (this takes ~30s on first run)...") | |
| try: | |
| _vision_model, _vision_processor = ml.load_vision_encoder() | |
| logger.info("BLIP-2 vision model loaded — image content will be analysed.") | |
| except Exception as exc: | |
| logger.warning(f"BLIP-2 failed to load: {exc}. Reports will use form-based captions only.") | |
| _models_loaded = True | |
| logger.info("All models loaded.") | |
| async def startup_event(): | |
| try: | |
| _ensure_models_loaded() | |
| except Exception as exc: | |
| logger.error(f"Startup model loading failed: {exc}") | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| # Health & Info | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| def health_check(): | |
| import torch | |
| return { | |
| "status": "healthy", | |
| "models_loaded": _models_loaded, | |
| "vision_enabled": ENABLE_VISION_CAPTION, | |
| "vision_model_loaded": _vision_model is not None, | |
| "device": "cuda" if torch.cuda.is_available() else "cpu", | |
| "timestamp": datetime.utcnow().isoformat(), | |
| } | |
| def list_models(): | |
| import torch | |
| return { | |
| "slm_model": os.getenv("SLM_MODEL_ID", "microsoft/phi-2"), | |
| "vision_model": os.getenv("VISION_MODEL_ID", "Salesforce/blip2-opt-2.7b"), | |
| "vision_enabled": ENABLE_VISION_CAPTION, | |
| "vision_loaded": _vision_model is not None, | |
| "device": "cuda" if torch.cuda.is_available() else "cpu", | |
| "models_loaded": _models_loaded, | |
| } | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| # Main report generation endpoint | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| async def generate_report( | |
| image: UploadFile = File(...), | |
| patient_name: str = Form(""), | |
| patient_age: int = Form(...), | |
| patient_gender: str = Form(...), | |
| symptoms: str = Form(...), | |
| clinical_history: str = Form(...), | |
| modality: str = Form(...), | |
| body_region: str = Form(...), | |
| view: Optional[str] = Form(""), | |
| ): | |
| # 1. Read image bytes | |
| file_bytes = await image.read() | |
| filename = image.filename or "upload.jpg" | |
| logger.info(f"Image received: {filename} ({len(file_bytes)/1024:.1f} KB)") | |
| # 2. Preprocess image | |
| try: | |
| pp = _pp() | |
| preprocessor = pp.RadiologyPreprocessor( | |
| modality=modality.lower(), | |
| body_region=body_region.lower(), | |
| ) | |
| pil_image = preprocessor.preprocess_pil(file_bytes, filename) | |
| except Exception as exc: | |
| logger.error(f"Preprocessing error: {exc}") | |
| raise HTTPException(status_code=422, detail=f"Image preprocessing failed: {exc}") | |
| # 3. Load models | |
| try: | |
| _ensure_models_loaded() | |
| except Exception as exc: | |
| raise HTTPException(status_code=503, detail=f"Model loading failed: {exc}") | |
| # 4. Generate image caption using BLIP-2 (reads actual image pixels) | |
| image_caption = ( | |
| f"{modality} image of the {body_region}. " | |
| f"Patient presents with {symptoms}." | |
| ) | |
| if ENABLE_VISION_CAPTION and _vision_model is not None: | |
| try: | |
| ml = _ml() | |
| _, blip_caption = ml.extract_visual_features( | |
| pil_image, _vision_model, _vision_processor | |
| ) | |
| if blip_caption and len(blip_caption) > 20: | |
| image_caption = blip_caption | |
| logger.info(f"BLIP-2 caption used: {image_caption[:120]}") | |
| except Exception as exc: | |
| logger.warning(f"Vision captioning failed: {exc} — using fallback caption") | |
| # 5. Generate report using Phi-2 | |
| try: | |
| rg = _rg() | |
| patient = rg.PatientInfo( | |
| age=patient_age, gender=patient_gender, | |
| symptoms=symptoms, clinical_history=clinical_history, | |
| ) | |
| imaging = rg.ImagingInfo( | |
| modality=modality, body_region=body_region, view=view or "" | |
| ) | |
| generator = rg.ReportGenerator( | |
| slm_model=_slm_model, slm_tokenizer=_slm_tokenizer | |
| ) | |
| report = generator.generate( | |
| patient=patient, imaging=imaging, image_caption=image_caption | |
| ) | |
| logger.info(f"Report generated in {report.generation_time_seconds:.1f}s") | |
| except Exception as exc: | |
| logger.error(f"Report generation error: {exc}") | |
| raise HTTPException(status_code=500, detail=f"Report generation failed: {exc}") | |
| # 6. Validate | |
| vm = _vm() | |
| validation = vm.validate_report(report) | |
| report.confidence_score = validation.confidence_score | |
| # 7. Medical coding (UMLS / ICD-10 / SNOMED CT) | |
| mc = _mc() | |
| coding_result = mc.extract_medical_codes( | |
| report_text=report.raw_text, | |
| sections={ | |
| "findings": report.findings, | |
| "impression": report.impression, | |
| "recommendations": report.recommendations, | |
| }, | |
| ) | |
| # 8. Build and store response | |
| report_id = str(uuid.uuid4()) | |
| # Clean patient name | |
| patient_name_clean = patient_name.strip() if patient_name.strip() else "Unknown Patient" | |
| payload = { | |
| "report_id": report_id, | |
| "timestamp": datetime.utcnow().isoformat(), | |
| "image_caption": image_caption, | |
| "patient_name": patient_name_clean, | |
| "patient": { | |
| "name": patient_name_clean, | |
| "age": patient_age, | |
| "gender": patient_gender, | |
| "symptoms": symptoms, | |
| "clinical_history": clinical_history, | |
| }, | |
| "imaging": { | |
| "modality": modality, | |
| "body_region": body_region, | |
| "view": view, | |
| }, | |
| "report": report.to_dict(), | |
| "report_text": report.to_text(), | |
| "validation": validation.to_dict(), | |
| "medical_coding": coding_result.to_dict(), | |
| } | |
| _report_store[report_id] = payload | |
| # Save to SQLite database | |
| try: | |
| db.save_report(payload, patient_name_clean) | |
| except Exception as db_exc: | |
| logger.warning(f"DB save failed (non-critical): {db_exc}") | |
| logger.info( | |
| f"Report stored: {report_id[:8]} | " | |
| f"confidence: {report.confidence_score:.2f} | " | |
| f"codes: {len(coding_result.codes)} | " | |
| f"diagnosis: {coding_result.primary_diagnosis_label}" | |
| ) | |
| return JSONResponse(content=payload) | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| # Report retrieval & export | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| def get_report(report_id: str): | |
| if report_id not in _report_store: | |
| raise HTTPException(status_code=404, detail="Report not found.") | |
| return JSONResponse(content=_report_store[report_id]) | |
| def export_json(report_id: str): | |
| if report_id not in _report_store: | |
| raise HTTPException(status_code=404, detail="Report not found.") | |
| return Response( | |
| content=json.dumps(_report_store[report_id], indent=2).encode(), | |
| media_type="application/json", | |
| headers={"Content-Disposition": f'attachment; filename="report_{report_id[:8]}.json"'}, | |
| ) | |
| def export_pdf(report_id: str): | |
| if report_id not in _report_store: | |
| raise HTTPException(status_code=404, detail="Report not found.") | |
| try: | |
| pdf_bytes = _generate_pdf(_report_store[report_id]) | |
| except Exception as exc: | |
| raise HTTPException(status_code=500, detail=f"PDF generation failed: {exc}") | |
| return Response( | |
| content=pdf_bytes, | |
| media_type="application/pdf", | |
| headers={"Content-Disposition": f'attachment; filename="report_{report_id[:8]}.pdf"'}, | |
| ) | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| # PDF generator | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| # Database endpoints | |
| # ═══════════════════════════════════════════════════════════════════════════ | |
| def list_all_reports(limit: int = 50, offset: int = 0): | |
| """List all stored reports from database.""" | |
| try: | |
| reports = db.get_all_reports(limit=limit, offset=offset) | |
| return JSONResponse(content={"reports": reports, "count": len(reports)}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| def get_patient_reports(patient_name: str): | |
| """Get all reports for a specific patient.""" | |
| try: | |
| reports = db.get_reports_by_patient(patient_name) | |
| return JSONResponse(content={"patient": patient_name, | |
| "reports": reports, | |
| "count": len(reports)}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| def search_patients(q: str = ""): | |
| """Search patients by name.""" | |
| try: | |
| results = db.search_patients(q) | |
| return JSONResponse(content={"results": results}) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| def get_stats(): | |
| """Get database statistics.""" | |
| try: | |
| return JSONResponse(content=db.get_db_stats()) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| def delete_report(report_id: str): | |
| """Delete a report by ID.""" | |
| try: | |
| success = db.delete_report(report_id) | |
| if not success: | |
| raise HTTPException(status_code=404, detail="Report not found.") | |
| return JSONResponse(content={"deleted": report_id}) | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| def _generate_pdf(payload: dict) -> bytes: | |
| from reportlab.lib import colors | |
| from reportlab.lib.enums import TA_CENTER, TA_JUSTIFY | |
| from reportlab.lib.pagesizes import A4 | |
| from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet | |
| from reportlab.lib.units import cm | |
| from reportlab.platypus import (HRFlowable, Paragraph, SimpleDocTemplate, | |
| Spacer, Table, TableStyle) | |
| buf = io.BytesIO() | |
| doc = SimpleDocTemplate(buf, pagesize=A4, | |
| rightMargin=2*cm, leftMargin=2*cm, | |
| topMargin=2.5*cm, bottomMargin=2*cm) | |
| styles = getSampleStyleSheet() | |
| BLUE = colors.HexColor("#1a4a8a") | |
| LIGHT = colors.HexColor("#e8f0fb") | |
| story = [] | |
| story.append(Paragraph("RADIOLOGY REPORT", | |
| ParagraphStyle("T", parent=styles["Title"], textColor=BLUE, | |
| fontSize=20, spaceAfter=4, fontName="Helvetica-Bold"))) | |
| report = payload.get("report", {}) | |
| patient = payload.get("patient", {}) | |
| imaging = payload.get("imaging", {}) | |
| validation = payload.get("validation", {}) | |
| report_id = payload.get("report_id", "N/A") | |
| timestamp = payload.get("timestamp", "") | |
| story.append(Paragraph( | |
| f"AI-Generated Report | {timestamp[:10]} | ID: {report_id[:8]}", | |
| ParagraphStyle("S", parent=styles["Normal"], textColor=colors.HexColor("#555"), | |
| fontSize=10, spaceAfter=12, alignment=TA_CENTER))) | |
| story.append(HRFlowable(width="100%", thickness=2, color=BLUE)) | |
| story.append(Spacer(1, 0.4*cm)) | |
| cell_style = ParagraphStyle("cell", parent=styles["Normal"], fontSize=9, | |
| leading=12) | |
| lbl_style = ParagraphStyle("lbl", parent=styles["Normal"], fontSize=9, | |
| leading=12, fontName="Helvetica-Bold") | |
| symptoms_text = str(patient.get("symptoms", "—")) | |
| tbl = Table([ | |
| [Paragraph("Patient Name", lbl_style), Paragraph(str(payload.get("patient_name","—")), cell_style), | |
| Paragraph("Report ID", lbl_style), Paragraph(str(report_id[:8]), cell_style)], | |
| [Paragraph("Patient Age", lbl_style), Paragraph(str(patient.get("age","—")), cell_style), | |
| Paragraph("Modality", lbl_style), Paragraph(str(imaging.get("modality","—")), cell_style)], | |
| [Paragraph("Gender", lbl_style), Paragraph(str(patient.get("gender","—")), cell_style), | |
| Paragraph("Body Region", lbl_style), Paragraph(str(imaging.get("body_region","—")), cell_style)], | |
| [Paragraph("View", lbl_style), Paragraph(str(imaging.get("view","—") or "—"), cell_style), | |
| Paragraph("", cell_style), Paragraph("", cell_style)], | |
| [Paragraph("Symptoms", lbl_style), Paragraph(symptoms_text, cell_style), | |
| Paragraph("", cell_style), Paragraph("", cell_style)], | |
| ], colWidths=[3*cm, 7*cm, 3*cm, 4*cm]) | |
| tbl.setStyle(TableStyle([ | |
| ("FONTSIZE", (0,0),(-1,-1), 9), | |
| ("BACKGROUND", (0,0),(-1,0), LIGHT), | |
| ("BACKGROUND", (0,2),(-1,2), LIGHT), | |
| ("GRID", (0,0),(-1,-1), 0.5, colors.HexColor("#ccc")), | |
| ("VALIGN", (0,0),(-1,-1), "TOP"), | |
| ("TOPPADDING", (0,0),(-1,-1), 5), | |
| ("BOTTOMPADDING", (0,0),(-1,-1), 5), | |
| ("SPAN", (1,2),(3,2)), | |
| ("SPAN", (1,3),(3,3)), | |
| ])) | |
| story.append(tbl) | |
| story.append(Spacer(1, 0.5*cm)) | |
| story.append(HRFlowable(width="100%", thickness=0.5, color=colors.HexColor("#aaa"))) | |
| hdr = ParagraphStyle("H", parent=styles["Heading2"], textColor=BLUE, | |
| fontSize=11, fontName="Helvetica-Bold", spaceBefore=12, spaceAfter=4) | |
| body = ParagraphStyle("B", parent=styles["Normal"], fontSize=10, | |
| leading=15, spaceAfter=6, alignment=TA_JUSTIFY) | |
| for header, key in [ | |
| ("Examination Details", "examination_details"), | |
| ("Clinical Information", "clinical_information"), | |
| ("Findings", "findings"), | |
| ("Impression", "impression"), | |
| ("Recommendations", "recommendations"), | |
| ]: | |
| story.append(Paragraph(header.upper(), hdr)) | |
| story.append(Paragraph(str(report.get(key,"—") or "—"), body)) | |
| story.append(Spacer(1, 0.5*cm)) | |
| story.append(HRFlowable(width="100%", thickness=1, color=BLUE)) | |
| confidence = validation.get("confidence_score", 0) | |
| story.append(Paragraph( | |
| f"CONFIDENCE SCORE: {confidence:.2f} / 1.00", | |
| ParagraphStyle("C", parent=styles["Normal"], fontSize=12, | |
| fontName="Helvetica-Bold", textColor=BLUE, | |
| alignment=TA_CENTER, spaceBefore=10))) | |
| story.append(Spacer(1, 1*cm)) | |
| story.append(HRFlowable(width="100%", thickness=0.5, color=colors.HexColor("#aaa"))) | |
| story.append(Paragraph( | |
| "This report was generated by an AI system. " | |
| "Clinical correlation by a licensed radiologist is mandatory.", | |
| ParagraphStyle("D", parent=styles["Normal"], fontSize=8, | |
| textColor=colors.HexColor("#888"), | |
| alignment=TA_CENTER, spaceBefore=6))) | |
| doc.build(story) | |
| return buf.getvalue() | |