Spaces:
Sleeping
Sleeping
File size: 2,657 Bytes
148cade e9e9fd2 148cade e9e9fd2 148cade e9e9fd2 148cade 56082d2 148cade 56082d2 148cade e9e9fd2 148cade 56082d2 148cade 9e43c15 148cade 9e43c15 148cade 9e43c15 148cade | 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 | import os
import tempfile
import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
# Import your prediction functions
from prediction import predict_eeg_recording, predict_ensemble_eeg_recording
app = FastAPI(title="EEG Epilepsy Prediction API")
@app.get("/", tags=["Introduction Endpoints"])
async def index():
"""
Simply returns a welcome message!
"""
message = (
"Hello world! Welcome to the EEG Epilepsy Prediction API. "
"Submit an EEG recording EDF file to the `/predict` endpoint to receive a prediction."
)
return message
@app.post("/predict", tags=["Machine Learning"])
async def predict_endpoint(
file: UploadFile = File(...),
model_choice: str = "2DCNN",
ensemble_method: str = None
):
"""
Query parameters:
- model_choice: Choose one model among "2DCNN", "EEGNet", "EpilepsyNet", or "ensemble".
- ensemble_method: (Optional, required if model_choice is "ensemble")
The ensemble method to use ("average" or "voting").
"""
print("Saving uploaded file as temporary file...")
try:
suffix = os.path.splitext(file.filename)[1]
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
tmp.write(await file.read())
tmp_path = tmp.name
except Exception as e:
raise HTTPException(status_code=500, detail="Error saving temporary file")
print("Performing prediction using model_choice =", model_choice)
try:
if model_choice.lower() == "ensemble":
if ensemble_method is None:
raise HTTPException(status_code=400, detail="ensemble_method must be specified when using ensemble model_choice")
pred_label, mean_prob, segment_probs = predict_ensemble_eeg_recording(
tmp_path, ensemble_method=ensemble_method, threshold=0.5
)
else:
pred_label, mean_prob, segment_probs = predict_eeg_recording(
tmp_path, model_name=model_choice, threshold=0.5
)
except Exception as e:
os.remove(tmp_path)
raise HTTPException(status_code=400, detail=f"Prediction failed: {e}")
os.remove(tmp_path)
response = {
"prediction": "epilepsy" if pred_label == 1 else "no epilepsy",
"mean_probability": float(mean_prob),
"segment_probabilities": [float(p) for p in segment_probs]
}
print("Prediction complete, returning response...")
return JSONResponse(content=response)
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)
|