fix visualisations + add heatmap

app.py

@@ -1,22 +1,28 @@
 import gradio as gr
 import numpy as np
+import matplotlib
 import matplotlib.pyplot as plt
 from sklearn.datasets import load_iris
 from sklearn.ensemble import GradientBoostingClassifier
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import accuracy_score, confusion_matrix
 
+# This line ensures Matplotlib doesn't try to open windows in certain environments:
+matplotlib.use('Agg')
+
+# Load the Iris dataset
 iris = load_iris()
 X, y = iris.data, iris.target
 feature_names = iris.feature_names
 class_names = iris.target_names
 
+# Train/test split
 X_train, X_test, y_train, y_test = train_test_split(
     X, y, test_size=0.3, random_state=42
 )
 
 def train_and_evaluate(learning_rate, n_estimators, max_depth):
-    # Train model
+    # Train the model
     clf = GradientBoostingClassifier(
         learning_rate=learning_rate,
         n_estimators=n_estimators,
@@ -25,29 +31,52 @@ def train_and_evaluate(learning_rate, n_estimators, max_depth):
     )
     clf.fit(X_train, y_train)
 
-    # Predict and compute metrics
+    # Predict on test set
     y_pred = clf.predict(X_test)
+
+    # Calculate accuracy
     accuracy = accuracy_score(y_test, y_pred)
+
+    # Calculate confusion matrix
     cm = confusion_matrix(y_test, y_pred)
 
-    # Convert confusion matrix to readable string
-    cm_display = "\n".join([str(row) for row in cm])
+    # Create a single figure with 2 subplots
+    fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(10, 4))
 
-    # Create a feature importance bar chart
+    # --- Subplot 1: Feature Importances ---
     importances = clf.feature_importances_
-    fig, ax = plt.subplots()
-    ax.barh(range(len(feature_names)), importances, color='skyblue')
-    ax.set_yticks(range(len(feature_names)))
-    ax.set_yticklabels(feature_names)
-    ax.set_xlabel("Importance")
-    ax.set_title("Feature Importances (Gradient Boosting)")
-
-    # Convert the Matplotlib figure to a Gradio-readable format
-    # (returns a temporary .png file path)
-    return (
-        f"Accuracy: {accuracy:.3f}\nConfusion Matrix:\n{cm_display}",
-        fig
-    )
+    axs[0].barh(range(len(feature_names)), importances, color='skyblue')
+    axs[0].set_yticks(range(len(feature_names)))
+    axs[0].set_yticklabels(feature_names)
+    axs[0].set_xlabel("Importance")
+    axs[0].set_title("Feature Importances")
+
+    # --- Subplot 2: Confusion Matrix Heatmap ---
+    im = axs[1].imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
+    axs[1].set_title("Confusion Matrix")
+    # Add colorbar
+    cbar = fig.colorbar(im, ax=axs[1])
+    # Tick marks for x/y axes
+    axs[1].set_xticks(range(len(class_names)))
+    axs[1].set_yticks(range(len(class_names)))
+    axs[1].set_xticklabels(class_names, rotation=45, ha="right")
+    axs[1].set_yticklabels(class_names)
+    axs[1].set_ylabel('True Label')
+    axs[1].set_xlabel('Predicted Label')
+
+    # Write the counts in each cell
+    thresh = cm.max() / 2.0
+    for i in range(cm.shape[0]):
+        for j in range(cm.shape[1]):
+            color = "white" if cm[i, j] > thresh else "black"
+            axs[1].text(j, i, format(cm[i, j], "d"),
+                        ha="center", va="center", color=color)
+
+    plt.tight_layout()
+
+    # Return textual results + the figure
+    results_text = f"Accuracy: {accuracy:.3f}"
+    return results_text, fig
 
 def predict_species(sepal_length, sepal_width, petal_length, petal_width,
                     learning_rate, n_estimators, max_depth):
@@ -65,13 +94,14 @@ def predict_species(sepal_length, sepal_width, petal_length, petal_width,
 with gr.Blocks() as demo:
     with gr.Tab("Train & Evaluate"):
         gr.Markdown("## Train a GradientBoostingClassifier on the Iris dataset")
+
         learning_rate_slider = gr.Slider(0.01, 1.0, value=0.1, step=0.01, label="learning_rate")
         n_estimators_slider = gr.Slider(50, 300, value=100, step=50, label="n_estimators")
         max_depth_slider = gr.Slider(1, 10, value=3, step=1, label="max_depth")
 
         train_button = gr.Button("Train & Evaluate")
         output_text = gr.Textbox(label="Results")
-        output_plot = gr.Plot(label="Feature Importance")
+        output_plot = gr.Plot(label="Feature Importances & Confusion Matrix")
 
         train_button.click(
             fn=train_and_evaluate,
@@ -81,14 +111,15 @@ with gr.Blocks() as demo:
 
     with gr.Tab("Predict"):
         gr.Markdown("## Predict Iris Species with GradientBoostingClassifier")
+
         sepal_length_input = gr.Number(value=5.1, label=feature_names[0])
-        sepal_width_input = gr.Number(value=3.5, label=feature_names[1])
+        sepal_width_input  = gr.Number(value=3.5, label=feature_names[1])
         petal_length_input = gr.Number(value=1.4, label=feature_names[2])
-        petal_width_input = gr.Number(value=0.2, label=feature_names[3])
+        petal_width_input  = gr.Number(value=0.2, label=feature_names[3])
 
-        learning_rate_slider2 = gr.Slider(0.01, 1.0, value=0.1, step=0.01, label="learning_rate")
-        n_estimators_slider2 = gr.Slider(50, 300, value=100, step=50, label="n_estimators")
-        max_depth_slider2 = gr.Slider(1, 10, value=3, step=1, label="max_depth")
+        learning_rate_slider2   = gr.Slider(0.01, 1.0, value=0.1, step=0.01, label="learning_rate")
+        n_estimators_slider2    = gr.Slider(50, 300, value=100, step=50, label="n_estimators")
+        max_depth_slider2       = gr.Slider(1, 10, value=3, step=1, label="max_depth")
 
         predict_button = gr.Button("Predict")
         prediction_text = gr.Textbox(label="Prediction")