Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,194 +1,135 @@
 # mushroom_app.py
 import streamlit as st
 import pandas as pd
-import numpy as np
-import joblib
 import requests
 from io import StringIO
 from sklearn.model_selection import train_test_split
-from sklearn.preprocessing import LabelEncoder
 from sklearn.ensemble import RandomForestClassifier
-from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
-import matplotlib.pyplot as plt
-import seaborn as sns
+from sklearn.preprocessing import LabelEncoder
+import joblib
+import os
 
 # -------------------------
-# Auto Download Mushroom Dataset
+# Auto Load Dataset
 # -------------------------
-@st.cache_data(show_spinner="Downloading Mushroom Dataset from UCI...")
-def load_mushroom_data():
+@st.cache_data
+def load_data():
     url = "https://archive.ics.uci.edu/ml/machine-learning-databases/mushroom/agaricus-lepiota.data"
     response = requests.get(url)
-    if response.status_code == 200:
-        # Column names as per UCI description
-        columns = [
-            'class', 'cap-shape', 'cap-surface', 'cap-color', 'bruises', 'odor',
-            'gill-attachment', 'gill-spacing', 'gill-size', 'gill-color',
-            'stalk-shape', 'stalk-root', 'stalk-surface-above-ring',
-            'stalk-surface-below-ring', 'stalk-color-above-ring',
-            'stalk-color-below-ring', 'veil-type', 'veil-color', 'ring-number',
-            'ring-type', 'spore-print-color', 'population', 'habitat'
-        ]
-        df = pd.read_csv(StringIO(response.text), header=None, names=columns)
-        return df
-    else:
-        st.error("Failed to download dataset.")
-        return None
-
-# Load data
-df = load_mushroom_data()
-if df is None:
-    st.stop()
-
-st.set_page_config(page_title="Mushroom Classification", layout="centered")
+    response.raise_for_status()
+    
+    columns = [
+        'class', 'cap_shape', 'cap_surface', 'cap_color', 'bruises', 'odor',
+        'gill_attachment', 'gill_spacing', 'gill_size', 'gill_color',
+        'stalk_shape', 'stalk_root', 'stalk_surface_above_ring',
+        'stalk_surface_below_ring', 'stalk_color_above_ring',
+        'stalk_color_below_ring', 'veil_type', 'veil_color', 'ring_number',
+        'ring_type', 'spore_print_color', 'population', 'habitat'
+    ]
+    df = pd.read_csv(StringIO(response.text), header=None, names=columns)
+    return df
+
+df = load_data()
+
+st.set_page_config(page_title="Mushroom Classifier", layout="centered")
 st.title("Mushroom Classification")
-st.markdown("### Is it *Edible* or *Poisonous*?")
-st.success(f"Dataset loaded: {df.shape[0]:,} mushrooms | {df.shape[1]} features")
-
-# Show sample
-st.write("First 10 mushrooms:")
-st.dataframe(df.head(10), use_container_width=True)
+st.markdown("### *Edible* or *Poisonous*? Let AI decide!")
 
-# Mapping
-class_map = {'e': 'Edible', 'p': 'Poisonous'}
-st.write("*Target Distribution:*")
+# Show data
+st.write(f"{len(df):,} mushrooms loaded**")
 col1, col2 = st.columns(2)
 with col1:
-    st.metric("Edible (e)", df['class'].value_counts().get('e', 0))
+    st.metric("Edible", len(df[df['class'] == 'e']))
 with col2:
-    st.metric("Poisonous (p)", df['class'].value_counts().get('p', 0))
+    st.metric("Poisonous", len(df[df['class'] == 'p']))
 
-# Encode categorical features
+# Preprocess
 @st.cache_data
-def preprocess_data(df):
+def preprocess(df):
     le_dict = {}
-    df_encoded = df.copy()
-    
-    for column in df.columns:
+    df_enc = df.copy()
+    for col in df.columns:
         le = LabelEncoder()
-        df_encoded[column] = le.fit_transform(df[column])
-        le_dict[column] = le
-    
-    X = df_encoded.drop('class', axis=1)
-    y = df_encoded['class']
-    
-    return X, y, le_dict, df_encoded
-
-X, y, label_encoders, df_encoded = preprocess_data(df)
+        df_enc[col] = le.fit_transform(df[col])
+        le_dict[col] = le
+    X = df_enc.drop('class', axis=1)
+    y = df_enc['class']
+    return X, y, le_dict
 
-# Train-test split
+X, y, encoders = preprocess(df)
 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
 
-# -------------------------
 # Train Model
-# -------------------------
-st.header("Train Random Forest Model (Best for this dataset)")
-if st.button("Train Model Now", type="primary"):
-    with st.spinner("Training Random Forest... (this dataset is 100% separable!)"):
+if st.button("Train Model (100% Accuracy Possible)", type="primary"):
+    with st.spinner("Training Random Forest..."):
         model = RandomForestClassifier(n_estimators=100, random_state=42)
         model.fit(X_train, y_train)
         
-        y_pred = model.predict(X_test)
-        acc = accuracy_score(y_test, y_pred)
-        
-        st.success(f"Model Trained! Accuracy: {acc*100:.2f}%")
-        
-        if acc >= 0.999:
+        acc = model.score(X_test, y_test)
+        st.success(f"Model Trained! Accuracy: {acc:.4%}")
+        if acc == 1.0:
             st.balloons()
-            st.markdown("### PERFECT CLASSIFICATION! (100% Accuracy)")
-
-        # Confusion Matrix
-        cm = confusion_matrix(y_test, y_pred)
-        fig, ax = plt.subplots(figsize=(6,5))
-        sns.heatmap(cm, annot=True, fmt='d', cmap="Greens", ax=ax,
-                    xticklabels=['Edible', 'Poisonous'],
-                    yticklabels=['Edible', 'Poisonous'])
-        ax.set_xlabel("Predicted")
-        ax.set_ylabel("Actual")
-        ax.set_title("Confusion Matrix")
-        st.pyplot(fig)
-
-        # Classification Report
-        report = classification_report(y_test, y_pred, target_names=['Edible', 'Poisonous'], output_dict=True)
-        report_df = pd.DataFrame(report).transpose()
-        st.write("### Classification Report")
-        st.dataframe(report_df.style.format("{:.3f}"))
+            st.markdown("*PERFECT CLASSIFICATION!*")
 
         # Save model
-        joblib.dump({
-            "model": model,
-            "label_encoders": label_encoders,
-            "features": X.columns.tolist()
-        }, "mushroom_model.pkl")
-        
-        with open("mushroom_model.pkl", "rb") as f:
-            st.download_button("Download Model (.pkl)", f, "mushroom_model.pkl")
+        joblib.dump({"model": model, "encoders": encoders}, "model.pkl")
+        st.session_state.model = model
+        st.session_state.encoders = encoders
 
-# -------------------------
-# Single Prediction (Perfect & Safe)
-# -------------------------
-st.header("Will You Eat This Mushroom?")
-st.markdown("Select features of a mushroom to predict if it's *safe to eat*")
-
-# Load model if trained
+# Load model if exists
 model = None
-if 'model' in globals() or st.session_state.get("model_trained", False):
-    try:
-        loaded = joblib.load("mushroom_model.pkl")
-        model = loaded["model"]
-        label_encoders = loaded["label_encoders"]
-        st.session_state.model_trained = True
-    except:
-        pass
-
+encoders = None
+if os.path.exists("model.pkl"):
+    data = joblib.load("model.pkl")
+    model = data["model"]
+    encoders = data["encoders"]
+elif "model" in st.session_state:
+    model = st.session_state.model
+    encoders = st.session_state.encoders
+
+# Prediction UI
+st.header("Test a Mushroom")
 if model is None:
-    st.info("Train the model above first to make predictions!")
+    st.info("Click 'Train Model' above to enable predictions")
 else:
     cols = st.columns(3)
     inputs = {}
     
-    feature_descriptions = {
-        'cap-shape': ['bell', 'conical', 'flat', 'knobbed', 'sunken', 'convex'],
-        'cap-surface': ['fibrous', 'grooves', 'smooth', 'scaly'],
-        'cap-color': ['buff', 'cinnamon', 'red', 'gray', 'brown', 'pink', 'green', 'purple', 'white', 'yellow'],
-        'bruises': ['yes', 'no'],
-        'odor': ['almond', 'creosote', 'foul', 'anise', 'musty', 'none', 'pungent', 'spicy', 'fishy'],
-        'gill-color': ['buff', 'red', 'gray', 'chocolate', 'black', 'brown', 'orange', 'pink', 'green', 'purple', 'white', 'yellow'],
-        'stalk-shape': ['enlarging', 'tapering'],
-        'stalk-root': ['bulbous', 'club', 'equal', 'rooted', 'missing'],
-        'spore-print-color': ['black', 'brown', 'buff', 'chocolate', 'green', 'orange', 'purple', 'white', 'yellow'],
+    feature_options = {
+        'cap_shape': ['bell', 'conical', 'convex', 'flat', 'knobbed', 'sunken'],
+        'cap_surface': ['fibrous', 'grooves', 'scaly', 'smooth'],
+        'cap_color': ['brown', 'buff', 'cinnamon', 'gray', 'green', 'pink', 'purple', 'red', 'white', 'yellow'],
+        'bruises': ['bruises', 'no'],
+        'odor': ['almond', 'anise', 'creosote', 'fishy', 'foul', 'musty', 'none', 'pungent', 'spicy'],
+        'gill_color': ['black', 'brown', 'buff', 'chocolate', 'gray', 'green', 'orange', 'pink', 'purple', 'red', 'white', 'yellow'],
+        'stalk_shape': ['enlarging', 'tapering'],
+        'spore_print_color': ['black', 'brown', 'buff', 'chocolate', 'green', 'orange', 'purple', 'white', 'yellow'],
         'population': ['abundant', 'clustered', 'numerous', 'scattered', 'several', 'solitary'],
-        'habitat': ['woods', 'grasses', 'leaves', 'meadows', 'paths', 'urban', 'waste']
+        'habitat': ['grasses', 'leaves', 'meadows', 'paths', 'urban', 'waste', 'woods']
     }
 
-    selected = {}
     for i, col in enumerate(X.columns):
+        options = feature_options.get(col, list(encoders[col].classes_))
         with cols[i % 3]:
-            options = feature_descriptions.get(col, ['?'])
-            idx = st.selectbox(f"{col.replace('-', ' ').title()}", options, key=col)
-            # Map to encoded value
-            original_values = label_encoders[col].classes_
-            if idx in original_values:
-                encoded_val = label_encoders[col].transform([idx])[0]
-                selected[col] = encoded_val
+            val = st.selectbox(col.replace("_", " ").title(), options, key=col)
+            code = encoders[col].transform([val])[0]
+            inputs[col] = code
 
-    if st.button("Predict Edible or Poisonous?", type="secondary"):
-        input_vec = [selected[f] for f in X.columns]
-        input_df = pd.DataFrame([input_vec], columns=X.columns)
-        
-        pred = model.predict(input_df)[0]
-        prob = model.predict_proba(input_df)[0]
-        
-        if pred == label_encoders['class'].transform(['e'])[0]:
-            st.success("EDIBLE & SAFE TO EAT!")
+    if st.button("Is it Safe to Eat?", type="secondary"):
+        input_vec = [[inputs[col] for col in X.columns]]
+        pred = model.predict(input_vec)[0]
+        prob = model.predict_proba(input_vec)[0]
+
+        if encoders['class'].inverse_transform([pred])[0] == 'e':
+            st.success("EDIBLE – Safe to Eat!")
             st.balloons()
         else:
-            st.error("POISONOUS! DO NOT EAT!")
-            st.warning("This mushroom is highly toxic!")
+            st.error("POISONOUS – DO NOT EAT!")
+            st.warning("Highly Toxic!")
 
         col1, col2 = st.columns(2)
-        col1.metric("Edible Probability", f"{prob[0]:.1%}")
-        col2.metric("Poisonous Probability", f"{prob[1]:.1%}")
+        col1.metric("Edible", f"{prob[0]:.1%}")
+        col2.metric("Poisonous", f"{prob[1]:.1%}")
 
 st.markdown("---")
-st.caption("Mushroom Classification • UCI Dataset • 100% Accuracy Possible • Built with Streamlit")
+st.caption("Mushroom Classifier • UCI Dataset • 100% Deployable on Hugging Face Spaces")