Update app.py

app.py

@@ -2,7 +2,10 @@ import gradio as gr
 import pandas as pd
 from sklearn.model_selection import train_test_split
 from sklearn.ensemble import RandomForestClassifier
-from sklearn.metrics import classification_report
+from sklearn.metrics import classification_report, confusion_matrix
+import matplotlib.pyplot as plt
+import seaborn as sns
+import io
 
 def load_data(file):
     if file is None:
@@ -13,93 +16,96 @@ def load_data(file):
         else:
             df = pd.read_excel(file.name)
         columns = list(df.columns)
-        return df, columns, df.head(100)
-    except Exception:
+        return df, columns, df.head(100)  # Show first 100 rows as preview
+    except Exception as e:
         return None, [], pd.DataFrame()
 
-def generate_dynamic_help(report_dict):
+def plot_confusion_matrix(y_true, y_pred, labels):
+    cm = confusion_matrix(y_true, y_pred, labels=labels)
+    plt.figure(figsize=(6,5))
+    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
+                xticklabels=labels, yticklabels=labels)
+    plt.xlabel("Predicted")
+    plt.ylabel("Actual")
+    plt.title("Confusion Matrix")
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    plt.close()
+    buf.seek(0)
+    return buf
+
+def generate_dynamic_help(report):
+    # Simple example: check for precision or recall < 0.5 and suggest caution
+    lines = report.splitlines()
     help_lines = []
-    classes = [k for k in report_dict if k not in ('accuracy', 'macro avg', 'weighted avg')]
-    
-    # Find class with lowest recall and lowest precision
-    lowest_recall_class = min(classes, key=lambda c: report_dict[c]['recall'])
-    lowest_precision_class = min(classes, key=lambda c: report_dict[c]['precision'])
-    
-    # Overall accuracy
-    accuracy = report_dict.get('accuracy', None)
-    if accuracy is not None:
-        help_lines.append(f"**Overall accuracy** of the model is {accuracy:.2f}.")
-    
-    # Comment on recall
-    recall_val = report_dict[lowest_recall_class]['recall']
-    if recall_val < 0.5:
-        help_lines.append(f"Class '{lowest_recall_class}' has a low recall ({recall_val:.2f}), meaning many true instances of this class are missed.")
-    else:
-        help_lines.append(f"Class '{lowest_recall_class}' has the lowest recall ({recall_val:.2f}), but it's above 0.5, which is reasonable.")
-    
-    # Comment on precision
-    precision_val = report_dict[lowest_precision_class]['precision']
-    if precision_val < 0.5:
-        help_lines.append(f"Class '{lowest_precision_class}' has a low precision ({precision_val:.2f}), indicating many false positives.")
-    else:
-        help_lines.append(f"Class '{lowest_precision_class}' has the lowest precision ({precision_val:.2f}), which is acceptable.")
-    
-    # Warn about low support classes
-    low_support_classes = [c for c in classes if report_dict[c]['support'] < 10]
-    if low_support_classes:
-        help_lines.append(f"Note: Classes {low_support_classes} have very few samples (support < 10), which may affect metric reliability.")
-    
-    # General advice
-    help_lines.append("Consider collecting more data or tuning the model if some classes show poor performance.")
-    
+    for line in lines:
+        if line.strip() == "":
+            continue
+        parts = line.split()
+        if len(parts) >= 4 and parts[0] not in ("accuracy", "macro", "weighted"):
+            try:
+                precision = float(parts[1])
+                recall = float(parts[2])
+                f1 = float(parts[3])
+                cls = parts[0]
+                if precision < 0.5:
+                    help_lines.append(f"⚠️ Precision for class **{cls}** is low ({precision:.2f}). The model often misclassifies samples as this class.")
+                if recall < 0.5:
+                    help_lines.append(f"⚠️ Recall for class **{cls}** is low ({recall:.2f}). The model misses many samples of this class.")
+            except:
+                continue
+    if not help_lines:
+        return "✅ Model performance looks good across all classes."
     return "\n\n".join(help_lines)
 
 def train_model(df, target_col, feature_cols):
     if df is None or df.empty:
-        return "Please upload a valid dataset first.", ""
+        return "Please upload a valid dataset first.", "", None
     if target_col not in df.columns:
-        return "Target column not found in dataset.", ""
+        return "Target column not found in dataset.", "", None
     if not feature_cols:
-        return "Please select at least one feature column.", ""
+        return "Please select at least one feature column.", "", None
 
     df_clean = df[[target_col] + feature_cols].dropna()
     if df_clean.empty:
-        return "No data left after removing missing values.", ""
+        return "No data left after removing missing values.", "", None
 
     X = df_clean[feature_cols]
     y = df_clean[target_col]
 
     if y.nunique() < 2:
-        return "Target must have at least 2 classes.", ""
+        return "Target must have at least 2 classes.", "", None
 
     X_enc = pd.get_dummies(X)
 
     try:
         X_train, X_test, y_train, y_test = train_test_split(X_enc, y, test_size=0.2, random_state=42)
     except ValueError as e:
-        return f"Error splitting data: {e}", ""
+        return f"Error splitting data: {e}", "", None
 
     if X_train.shape[0] == 0 or X_test.shape[0] == 0:
-        return "Empty train or test set after splitting.", ""
+        return "Empty train or test set after splitting.", "", None
 
     model = RandomForestClassifier(random_state=42)
     model.fit(X_train, y_train)
     y_pred = model.predict(X_test)
 
-    report_dict = classification_report(y_test, y_pred, output_dict=True)
-    report_str = classification_report(y_test, y_pred)
-    help_str = generate_dynamic_help(report_dict)
+    report = classification_report(y_test, y_pred)
+    dynamic_help = generate_dynamic_help(report)
 
-    return report_str, help_str
+    labels = sorted(y_test.unique())
+    cm_buf = plot_confusion_matrix(y_test, y_pred, labels)
+
+    return report, dynamic_help, cm_buf
 
 def on_file_change(file):
     df, columns, preview = load_data(file)
     if df is None:
-        return None, gr.update(choices=[], value=None), gr.update(choices=[], value=[]), pd.DataFrame()
-    return df, gr.update(choices=columns, value=None), gr.update(choices=columns, value=[]), preview
+        return None, gr.Dropdown.update(choices=[], value=None), gr.CheckboxGroup.update(choices=[], value=[]), pd.DataFrame()
+    return df, gr.Dropdown.update(choices=columns, value=None), gr.CheckboxGroup.update(choices=columns, value=[]), preview
 
 with gr.Blocks() as demo:
-    gr.Markdown("# XLSX/CSV Classification App with Dynamic Help")
+    gr.Markdown("# XLSX/CSV Classification App with Table Preview and Visualization")
 
     df_state = gr.State(None)
 
@@ -111,9 +117,14 @@ with gr.Blocks() as demo:
         target_col = gr.Dropdown(label="Select Target Column", choices=[])
     with gr.Row():
         feature_cols = gr.CheckboxGroup(label="Select Feature Columns", choices=[])
-    train_btn = gr.Button("Train Model")
-    output_report = gr.Textbox(label="Classification Report", lines=10)
-    output_help = gr.Markdown(label="Dynamic Help")
+    with gr.Row():
+        train_btn = gr.Button("Train Model")
+    with gr.Row():
+        output_report = gr.Textbox(label="Classification Report", lines=10)
+    with gr.Row():
+        output_help = gr.Markdown(label="Model Performance Help")
+    with gr.Row():
+        cm_image = gr.Image(label="Confusion Matrix")
 
     file_input.change(
         fn=on_file_change,
@@ -124,7 +135,7 @@ with gr.Blocks() as demo:
     train_btn.click(
         fn=train_model,
         inputs=[df_state, target_col, feature_cols],
-        outputs=[output_report, output_help]
+        outputs=[output_report, output_help, cm_image]
     )
 
 demo.launch()