Update src/app.py

src/app.py

@@ -1,78 +1,86 @@
 import io
 from pathlib import Path
-
 import streamlit as st
 from fastai.vision.all import load_learner, PILImage
 
-
+# ---------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------
 MODEL_PATH = Path("models/pokemon_gen9_classifier_resnet101_after_cleaning.pkl")
 
-
-@st.cache_resource
+# ---------------------------------------------------------------------
+# Load model (cached so it doesn't reload on every rerun)
+# ---------------------------------------------------------------------
+@st.cache_resource(show_spinner=False)
 def load_model():
     """Load and cache the FastAI learner. Returns None if model missing or incompatible."""
     if not MODEL_PATH.exists():
+        st.error(f"❌ Model file not found at: `{MODEL_PATH}`")
         return None
     try:
         learner = load_learner(MODEL_PATH)
         return learner
-    except RuntimeError as e:
-        if "deprecated in `fastai>=2.8.0`" in str(e):
-            st.error(
-                f"⚠️ Model incompatibility detected!\n\n"
-                f"Your model was exported with FastAI ≥2.8.0, but this app uses FastAI 2.7.12.\n\n"
-                f"**To fix this:**\n"
-                f"1. Re-export your model using FastAI 2.7.12:\n"
-                f"   - Downgrade: `pip install fastai==2.7.12 fastcore==1.7.9`\n"
-                f"   - Run: `learn.export('{MODEL_PATH}')`\n"
-                f"2. Place the new export at: `{MODEL_PATH}`\n"
-                f"3. Refresh this page.\n\n"
-                f"**Error details:** {str(e)}"
-            )
-            return None
-        raise
-
+    except Exception as e:
+        st.error(
+            f"⚠️ Failed to load model at `{MODEL_PATH}`.\n\n"
+            f"**Error:** {e}\n\n"
+            f"If this model was exported using a newer FastAI version, "
+            f"re-export it using FastAI 2.7.12:\n"
+            f"`pip install fastai==2.7.12 fastcore==1.5.29`\n"
+            f"`learn.export('{MODEL_PATH}')`"
+        )
+        return None
 
+# ---------------------------------------------------------------------
+# Prediction function
+# ---------------------------------------------------------------------
 def predict(learner, img_bytes: bytes):
+    """Run inference on uploaded image and return predictions."""
     img = PILImage.create(io.BytesIO(img_bytes))
     pred, pred_idx, probs = learner.predict(img)
     return pred, probs
 
-
+# ---------------------------------------------------------------------
+# Streamlit UI
+# ---------------------------------------------------------------------
 def main():
-    st.title("FastAI Image Classifier")
+    st.set_page_config(page_title="FastAI Image Classifier", layout="centered")
+    st.title("🧠 FastAI Image Classifier")
     st.write("Upload an image and the model will predict the class.")
 
     learner = load_model()
     if learner is None:
-        st.warning(
-            f"No model found at `{MODEL_PATH}`. See README.md for how to export and place your model."
-        )
+        st.warning(f"No model loaded. Make sure `{MODEL_PATH}` exists inside the container.")
+        return
 
-    uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
-    if uploaded_file is not None:
-        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+    uploaded_file = st.file_uploader("📤 Choose an image...", type=["png", "jpg", "jpeg"])
 
-        if learner is None:
-            st.error("Can't predict because the model is missing. Follow README.md to export your learner.")
-            return
+    if uploaded_file is not None:
+        st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
 
-        with st.spinner("Predicting..."):
-            pred, probs = predict(learner, uploaded_file.read())
+        with st.spinner("🔍 Predicting..."):
+            try:
+                pred, probs = predict(learner, uploaded_file.read())
+            except Exception as e:
+                st.error(f"Prediction failed: {e}")
+                return
 
-        st.success(f"Predicted: {pred}")
+        st.success(f"✅ **Predicted:** {pred}")
 
-        # Try to show top probabilities if available
+        # Show top 5 predictions (if available)
         try:
             vocab = learner.dls.vocab
-            probs_list = list(zip(map(str, vocab), map(float, probs)))
-            probs_list.sort(key=lambda x: x[1], reverse=True)
-            st.write("Top predictions:")
+            probs_list = sorted(
+                zip(map(str, vocab), map(float, probs)),
+                key=lambda x: x[1],
+                reverse=True
+            )
+            st.subheader("Top Predictions")
             for label, p in probs_list[:5]:
-                st.write(f"- {label}: {p:.4f}")
+                st.write(f"- **{label}**: {p*100:.2f}%")
         except Exception:
-            st.write("Probabilities unavailable or not applicable for this model.")
-
+            st.info("Probabilities unavailable or not applicable for this model.")
 
+# ---------------------------------------------------------------------
 if __name__ == "__main__":
     main()