Update app.py

app.py

@@ -3,14 +3,56 @@ from pathlib import Path
 import streamlit as st
 from fastai.vision.all import load_learner, PILImage
 
+# === CONFIG ===
+st.set_page_config(
+    page_title="FastAI Image Classifier",
+    page_icon="🎯",
+    layout="centered",
+    initial_sidebar_state="collapsed",
+)
 
-# ✅ Correct absolute path for Hugging Face Spaces
-MODEL_PATH = Path("models/pokemon_gen9_classifier_resnet101_after_cleaning.pkl") 
+MODEL_PATH = Path("models/pokemon_gen9_classifier_resnet101_after_cleaning.pkl")
+EXAMPLES_DIR = Path("examples")
 
+# === STYLES ===
+st.markdown("""
+    <style>
+        .block-container {
+            max-width: 750px;
+            margin: auto;
+            padding-top: 2rem;
+        }
+        .stButton button {
+            border-radius: 10px;
+            background-color: #4CAF50 !important;
+            color: white !important;
+            font-weight: 600;
+            padding: 0.6em 1.2em;
+            transition: all 0.2s ease-in-out;
+        }
+        .stButton button:hover {
+            background-color: #45a049 !important;
+            transform: scale(1.05);
+        }
+        h1, h2, h3 {
+            text-align: center;
+            color: #222;
+        }
+        .prediction-box {
+            background-color: #f7f7f7;
+            border-radius: 12px;
+            padding: 1rem;
+            margin-top: 1rem;
+            box-shadow: 0 2px 6px rgba(0,0,0,0.1);
+        }
+    </style>
+""", unsafe_allow_html=True)
 
+
+# === MODEL LOADING ===
 @st.cache_resource
 def load_model():
-    """Load and cache the FastAI learner. Returns None if model missing or incompatible."""
+    """Load and cache the FastAI learner."""
     if not MODEL_PATH.exists():
         st.error(f"❌ Model not found at {MODEL_PATH}")
         return None
@@ -29,33 +71,58 @@ def predict(learner, img_bytes: bytes):
     return pred, probs
 
 
+# === MAIN APP ===
 def main():
     st.title("🎯 FastAI Image Classifier")
-    st.write("Upload an image and the model will predict its class.")
+    st.write("Upload an image or try one of the examples below!")
 
     learner = load_model()
     if learner is None:
-        st.warning(
-            "Model not loaded. Please ensure the `.pkl` file is correctly placed under `models/` and committed with Git LFS."
-        )
+        st.warning("Please ensure your `.pkl` model is correctly placed under `models/` and committed with Git LFS.")
         st.stop()
 
-    uploaded_file = st.file_uploader("📤 Choose an image...", type=["png", "jpg", "jpeg"])
-    if uploaded_file is not None:
+    # Example images
+    example_images = list(EXAMPLES_DIR.glob("*"))
+    selected_example = None
+
+    if example_images:
+        st.subheader("✨ Try Example Images")
+        cols = st.columns(len(example_images))
+        for i, img_path in enumerate(example_images):
+            with cols[i]:
+                if st.button(img_path.stem.capitalize()):
+                    selected_example = img_path
+
+    uploaded_file = st.file_uploader("📤 Upload your own image", type=["png", "jpg", "jpeg"])
+
+    if selected_example:
+        img_bytes = open(selected_example, "rb").read()
+        st.image(selected_example, caption=f"Example: {selected_example.stem}", use_column_width=True)
+    elif uploaded_file:
+        img_bytes = uploaded_file.read()
         st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+    else:
+        st.stop()
+
+    st.markdown("---")
+
+    with st.spinner("🔍 Analyzing the image..."):
+        try:
+            pred, probs = predict(learner, img_bytes)
+            vocab = learner.dls.vocab
+            probs_list = sorted(zip(vocab, probs), key=lambda x: x[1], reverse=True)
+
+            # Display prediction
+            st.markdown(f"<div class='prediction-box'><h3>✅ Prediction: <span style='color:#4CAF50'>{pred}</span></h3></div>", unsafe_allow_html=True)
+            st.progress(float(probs.max()))
+
+            # Top 5 predictions
+            st.subheader("Top 5 Predictions")
+            for label, p in probs_list[:5]:
+                st.write(f"• **{label}** — {p:.2%}")
 
-        with st.spinner("Predicting..."):
-            try:
-                pred, probs = predict(learner, uploaded_file.read())
-                st.success(f"✅ Predicted: **{pred}**")
-                # Show top-5 predictions
-                vocab = learner.dls.vocab
-                probs_list = sorted(zip(vocab, probs), key=lambda x: x[1], reverse=True)
-                st.write("### Top Predictions:")
-                for label, p in probs_list[:5]:
-                    st.write(f"- {label}: {p:.4f}")
-            except Exception as e:
-                st.error(f"Error during prediction: {e}")
+        except Exception as e:
+            st.error(f"❌ Error during prediction:\n\n{e}")
 
 
 if __name__ == "__main__":