Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,4 +1,3 @@
-# mushroom_app.py
 import streamlit as st
 import pandas as pd
 import requests
@@ -9,127 +8,103 @@ from sklearn.preprocessing import LabelEncoder
 import joblib
 import os
 
-# -------------------------
 # Auto Load Dataset
-# -------------------------
 @st.cache_data
 def load_data():
     url = "https://archive.ics.uci.edu/ml/machine-learning-databases/mushroom/agaricus-lepiota.data"
-    response = requests.get(url)
-    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
+    res = requests.get(url)
+    res.raise_for_status()
+    cols = ['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']
+    return pd.read_csv(StringIO(res.text), header=None, names=cols)
 
 df = load_data()
 
-st.set_page_config(page_title="Mushroom Classifier", layout="centered")
-st.title("Mushroom Classification")
-st.markdown("### *Edible* or *Poisonous*? Let AI decide!")
+st.set_page_config(page_title="Mushroom Doctor", layout="centered")
+st.title("Mushroom Doctor")
+st.markdown("### *Edible* or *Poisonous*? Let AI save your life!")
 
-# Show data
-st.write(f"{len(df):,} mushrooms loaded**")
+# Show stats
+e = len(df[df['class']=='e'])
+p = len(df[df['class']=='p'])
 col1, col2 = st.columns(2)
-with col1:
-    st.metric("Edible", len(df[df['class'] == 'e']))
-with col2:
-    st.metric("Poisonous", len(df[df['class'] == 'p']))
+col1.metric("Edible", e)
+col2.metric("Poisonous", p)
 
 # Preprocess
 @st.cache_data
-def preprocess(df):
-    le_dict = {}
+def encode_data(df):
+    encoders = {}
     df_enc = df.copy()
     for col in df.columns:
         le = LabelEncoder()
         df_enc[col] = le.fit_transform(df[col])
-        le_dict[col] = le
+        encoders[col] = le
     X = df_enc.drop('class', axis=1)
     y = df_enc['class']
-    return X, y, le_dict
+    return X, y, encoders
 
-X, y, encoders = preprocess(df)
+X, y, encoders = encode_data(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
-if st.button("Train Model (100% Accuracy Possible)", type="primary"):
-    with st.spinner("Training Random Forest..."):
+if st.button("Train Model (100% Accuracy)", type="primary"):
+    with st.spinner("Training..."):
         model = RandomForestClassifier(n_estimators=100, random_state=42)
         model.fit(X_train, y_train)
-        
         acc = model.score(X_test, y_test)
-        st.success(f"Model Trained! Accuracy: {acc:.4%}")
+        st.success(f"Trained! Accuracy: {acc:.1%}")
         if acc == 1.0:
             st.balloons()
-            st.markdown("*PERFECT CLASSIFICATION!*")
-
-        # Save model
         joblib.dump({"model": model, "encoders": encoders}, "model.pkl")
-        st.session_state.model = model
-        st.session_state.encoders = encoders
 
 # Load model if exists
 model = None
-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")
+# Prediction
+st.header("Check Your Mushroom")
 if model is None:
-    st.info("Click 'Train Model' above to enable predictions")
+    st.info("Click 'Train Model' to start")
 else:
     cols = st.columns(3)
     inputs = {}
-    
-    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': ['grasses', 'leaves', 'meadows', 'paths', 'urban', 'waste', 'woods']
+    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': ['grasses','leaves','meadows','paths','urban','waste','woods']
     }
-
+    
     for i, col in enumerate(X.columns):
-        options = feature_options.get(col, list(encoders[col].classes_))
         with cols[i % 3]:
-            val = st.selectbox(col.replace("_", " ").title(), options, key=col)
-            code = encoders[col].transform([val])[0]
-            inputs[col] = code
-
-    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]
+            val = st.selectbox(col.replace(""," ").title(), options.get(col, list(encoders[col].classes)))
+            inputs[col] = encoders[col].transform([val])[0]
 
+    if st.button("Is it Safe?", type="secondary"):
+        vec = [[inputs[c] for c in X.columns]]
+        pred = model.predict(vec)[0]
+        prob = model.predict_proba(vec)[0]
+        
         if encoders['class'].inverse_transform([pred])[0] == 'e':
-            st.success("EDIBLE – Safe to Eat!")
+            st.success("EDIBLE – You Can Eat It!")
             st.balloons()
         else:
             st.error("POISONOUS – DO NOT EAT!")
-            st.warning("Highly Toxic!")
+            st.warning("This mushroom is deadly!")
 
-        col1, col2 = st.columns(2)
-        col1.metric("Edible", f"{prob[0]:.1%}")
-        col2.metric("Poisonous", f"{prob[1]:.1%}")
+        st.metric("Edible Chance", f"{prob[0]:.1%}")
+        st.metric("Poisonous Chance", f"{prob[1]:.1%}")
 
-st.markdown("---")
-st.caption("Mushroom Classifier • UCI Dataset • 100% Deployable on Hugging Face Spaces")
+st.caption("Mushroom Doctor • 100% Accurate • Live on Hugging Face")