Update app.py

app.py

@@ -32,11 +32,27 @@ app = FastAPI()
 @app.post("/predict")
 def predict(input_data: FatigueInput):
     try:
+        # Convert input to DataFrame
         input_dict = input_data.dict()
         input_df = pd.DataFrame([input_dict])[feature_cols]
+
+        # Scale the input
         scaled_input = scaler.transform(input_df)
+
+        # Get cluster prediction and probabilities
         cluster = gmm.predict(scaled_input)[0]
+        probs = gmm.predict_proba(scaled_input)[0]
+
+        # Map cluster to fatigue level
         fatigue_level = cluster_fatigue_map[str(cluster)]
-        return {"cluster": int(cluster), "fatigue_level": fatigue_level}
+
+        # Return more detailed prediction info
+        return {
+            "cluster": int(cluster),
+            "fatigue_level": fatigue_level,
+            "cluster_probabilities": {str(i): float(prob) for i, prob in enumerate(probs)},
+            "input_features": input_dict
+        }
+
     except Exception as e:
         return {"error": str(e)}