Update app.py

app.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
 import gradio as gr
+import matplotlib.pyplot as plt
 from sklearn.model_selection import train_test_split
 from sklearn.feature_selection import SelectKBest, f_classif
 from sklearn.preprocessing import StandardScaler
@@ -24,17 +25,12 @@ if not os.path.exists(MODEL_FILE):
         "contrast_intensity": np.random.uniform(50, 200, n_samples), # MRI contrast
     })
 
-    # Create balanced tumor labels
+    # Labels: tumor if size > 5cm or contrast > 120
     y = (
         (X["tumor_size"] > 5).astype(int) |
         (X["contrast_intensity"] > 120).astype(int)
     )
 
-    # Force balance: make half tumor, half no tumor
-    half = n_samples // 2
-    y[:half] = 0
-    y[half:] = 1
-
     print("Class distribution:", np.bincount(y))
 
     X_train, X_test, y_train, y_test = train_test_split(
@@ -57,14 +53,36 @@ if not os.path.exists(MODEL_FILE):
 # Load model
 model = joblib.load(MODEL_FILE)
 
-def predict_brain_tumor(tumor_size, texture_score, age, contrast_intensity):
+# Feature names (before selection)
+FEATURES = ["tumor_size", "texture_score", "age", "contrast_intensity"]
+
+def predict_and_explain(tumor_size, texture_score, age, contrast_intensity):
     input_data = np.array([[tumor_size, texture_score, age, contrast_intensity]])
     pred = model.predict(input_data)[0]
-    return "🧠 Tumor Detected" if pred == 1 else "✅ No Tumor Detected"
+
+    # Get feature importance from RandomForest
+    rf = model.named_steps["classifier"]
+    selector = model.named_steps["select"]
+
+    # Selected features
+    mask = selector.get_support()
+    selected_features = np.array(FEATURES)[mask]
+
+    importances = rf.feature_importances_
+
+    # Plot feature importances
+    fig, ax = plt.subplots(figsize=(6, 4))
+    ax.barh(selected_features, importances, color="royalblue")
+    ax.set_xlabel("Importance")
+    ax.set_title("Feature Importance for Tumor Detection")
+    plt.tight_layout()
+
+    result = "🧠 Tumor Detected" if pred == 1 else "✅ No Tumor Detected"
+    return result, fig
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Brain Tumor Detection")
-    gr.Markdown("Predict brain tumor presence using MRI-based features (simulated demo).")
+    gr.Markdown("# 🧠 Brain Tumor Detection with Feature Importance")
+    gr.Markdown("Predict brain tumor presence and view feature contributions (simulated demo).")
 
     with gr.Row():
         tumor_size = gr.Slider(1, 10, value=4, step=0.1, label="Tumor Size (cm)")
@@ -72,13 +90,14 @@ with gr.Blocks() as demo:
         age = gr.Slider(20, 80, value=40, step=1, label="Patient Age")
         contrast_intensity = gr.Slider(50, 200, value=100, step=1, label="MRI Contrast Intensity")
 
-    output = gr.Textbox(label="Prediction")
+    output_text = gr.Textbox(label="Prediction")
+    output_plot = gr.Plot(label="Feature Importance")
 
     predict_btn = gr.Button("Predict")
     predict_btn.click(
-        fn=predict_brain_tumor,
+        fn=predict_and_explain,
         inputs=[tumor_size, texture_score, age, contrast_intensity],
-        outputs=output,
+        outputs=[output_text, output_plot],
     )
 
 demo.launch()