Update app.py

app.py

@@ -11,14 +11,12 @@ 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
+    return {
+        "message": (
+            "Hello world! Welcome to the EEG Epilepsy Prediction API. "
+            "Submit an EEG recording EDF file to the `/predict` endpoint to receive a prediction."
+        )
+    }
 
 @app.post("/predict", tags=["Machine Learning"])
 async def predict_endpoint(
@@ -26,14 +24,6 @@ async def predict_endpoint(
     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]
@@ -41,23 +31,26 @@ async def predict_endpoint(
             tmp.write(await file.read())
             tmp_path = tmp.name
     except Exception as e:
-        raise HTTPException(status_code=500, detail="Error saving temporary file")
+        raise HTTPException(status_code=500, detail=f"Error saving temporary file: {e}")
     
     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 = predict_ensemble_eeg_recording(tmp_path, ensemble_method=ensemble_method, threshold=0.5)
+            pred_label, mean_prob, segment_probs = predict_ensemble_eeg_recording(
+                tmp_path, ensemble_method=ensemble_method, threshold=0.5
+            )
         else:
-            pred_label, mean_prob = predict_eeg_recording(tmp_path, model_name=model_choice, threshold=0.5)
+            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}")
+        raise HTTPException(status_code=500, detail=f"Prediction failed: {e}")
     
     os.remove(tmp_path)
-    
-    
+
     response = {
         "prediction": "epilepsy" if pred_label == 1 else "no epilepsy",
         "mean_probability": float(mean_prob),