Update app.py

app.py

@@ -5,7 +5,7 @@ from sklearn.datasets import load_iris
 from sklearn.ensemble import RandomForestClassifier
 
 app = Flask(__name__)
-CORS(app)  # Allow frontend calls from GitHub Pages
+CORS(app)
 
 # --- Train or load model
 try:
@@ -39,7 +39,13 @@ def home():
 }</pre>
     <p>Response:</p>
     <pre>{
-  "prediction": "setosa"
+  "prediction": "setosa",
+  "confidence": 0.98,
+  "probabilities": {
+      "setosa": 0.98,
+      "versicolor": 0.01,
+      "virginica": 0.01
+  }
 }</pre>
     """
 
@@ -54,15 +60,26 @@ def predict():
         data["petal_width"]
     ]
 
-    prediction = model.predict([features])[0]  # could be int or string
+    # Predict class and probabilities
+    prediction_idx = model.predict([features])[0]
+    probs = model.predict_proba([features])[0]
 
-    # If prediction is already a string (class label)
-    if isinstance(prediction, str):
-        return jsonify({"prediction": prediction})
-
-    # If prediction is numeric (0, 1, 2)
     target_names = load_iris().target_names
-    return jsonify({"prediction": target_names[int(prediction)]})
+    prediction_label = target_names[int(prediction_idx)]
+
+    # Build probability dict
+    probabilities = {
+        target_names[i]: float(probs[i])
+        for i in range(len(target_names))
+    }
+
+    confidence = float(max(probs))
+
+    return jsonify({
+        "prediction": prediction_label,
+        "confidence": confidence,
+        "probabilities": probabilities
+    })
 
 
 if __name__ == "__main__":