Spaces:
Sleeping
Sleeping
File size: 4,044 Bytes
5371cce d85bf0a 5371cce | 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 | import os
from fastapi import FastAPI, UploadFile, File, Form, Depends, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from sqlalchemy.orm import Session
from database import init_db, get_db, PredictionRecord
from model_loader import preload_all_models
from predict import predict_image
app = FastAPI(title="AI-Assisted Diagnostic Support API", version="1.0.0")
# Enable CORS for frontend integration
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # For local development simplicity
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Startup event: Initialize database and preload models
@app.on_event("startup")
def startup_event():
print("Starting up diagnostic API...")
print("Initializing SQLite Database...")
init_db()
# Preload models to avoid lag on first request
try:
preload_all_models()
print("All models preloaded successfully and ready.")
except Exception as e:
print(f"Warning: Failed to preload models during startup: {str(e)}")
@app.get("/health")
def health_check():
return {
"status": "healthy",
"models_loaded": True,
"supported_modules": ["blood_cell", "skin_lesion"]
}
@app.post("/predict")
async def run_predict(
file: UploadFile = File(...),
module_type: str = Form(...), # 'blood_cell' or 'skin_lesion'
db: Session = Depends(get_db)
):
if module_type not in ["blood_cell", "skin_lesion"]:
raise HTTPException(status_code=400, detail="Invalid module type. Choose 'blood_cell' or 'skin_lesion'.")
try:
# Read uploaded image bytes
image_bytes = await file.read()
# Run inference pipeline
label, confidence, class_probs, heatmap_b64, inf_time = predict_image(
image_bytes=image_bytes,
filename=file.filename,
module_type=module_type
)
# Save record in the SQLite database
record = PredictionRecord(
module=module_type,
filename=file.filename,
predicted_class=label,
confidence=confidence,
notes=f"Inference took {inf_time:.4f} seconds."
)
db.add(record)
db.commit()
db.refresh(record)
return {
"id": record.id,
"timestamp": record.timestamp,
"filename": record.filename,
"prediction": label,
"confidence": confidence,
"class_probabilities": class_probs,
"heatmap_base64": heatmap_b64,
"inference_time_sec": inf_time
}
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Inference execution failed: {str(e)}")
@app.get("/history")
def get_prediction_history(db: Session = Depends(get_db)):
records = db.query(PredictionRecord).order_by(PredictionRecord.timestamp.desc()).all()
# Format database records into JSON-friendly structures
history_list = []
for r in records:
history_list.append({
"id": r.id,
"timestamp": r.timestamp.isoformat() + "Z" if r.timestamp else None,
"module": "Blood Cell" if r.module == "blood_cell" else "Skin Lesion",
"filename": r.filename,
"prediction": r.predicted_class,
"confidence": r.confidence,
"notes": r.notes
})
return history_list
# Serve static files and index.html
static_dir = os.path.join(os.path.dirname(__file__), "static")
if os.path.exists(static_dir):
app.mount("/static", StaticFiles(directory=static_dir), name="static")
@app.get("/")
def read_index():
return FileResponse(os.path.join(static_dir, "index.html"))
if __name__ == "__main__":
import uvicorn
uvicorn.run("backend.app:app", host="0.0.0.0", port=8000, reload=True)
|