Update src/app.py

src/app.py

@@ -1,6 +1,4 @@
 import os
-from pathlib import Path
-
 import streamlit as st
 from streamlit_drawable_canvas import st_canvas
 import numpy as np
@@ -8,125 +6,91 @@ from tensorflow.keras.models import load_model
 from tensorflow.keras.preprocessing.image import img_to_array
 from PIL import Image
 
-# -------------------
-# Branding & Settings
-# -------------------
+# Reduce TensorFlow warnings
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
+
+# Branding Colors
 PRIMARY_COLOR = "#0F2C59"   # Procelevate navy
-ACCENT_COLOR  = "#FFD700"   # Gold highlight
-BG_COLOR      = "#F5F5F5"   # Light background
+ACCENT_COLOR = "#FFD700"    # Gold highlight
+BG_COLOR = "#F5F5F5"        # Light background
+
+# Load Model
+try:
+    model = load_model("src/quickdraw_model.keras")
+except Exception as e:
+    st.error(f"❌ Could not load model: {e}")
+    st.stop()
+
 CLASSES = ["cat", "dog", "house", "car", "tree", "airplane", "sun", "fish", "flower", "bird"]
 
+# Emoji map for results
+emoji_map = {
+    "cat": "🐱", "dog": "🐶", "car": "🚗", "tree": "🌳",
+    "house": "🏠", "sun": "☀️", "airplane": "✈️",
+    "fish": "🐟", "flower": "🌸", "bird": "🐦"
+}
+
+# Page Config
 st.set_page_config(page_title="Guess the Doodle - Procelevate", page_icon="✏️", layout="wide")
 
-# Sidebar
-st.sidebar.image("src/procelevate_logo.png")  # also works if copied to /app/src
+# Sidebar with logo
+st.sidebar.image("src/procelevate_logo.png")
 st.sidebar.markdown(
     f"<h3 style='color:{PRIMARY_COLOR};'>🚀 Procelevate Academy</h3>"
     "<p>Where AI meets fun & learning!</p>",
     unsafe_allow_html=True
 )
 
-# -------------------
-# Robust model loader
-# -------------------
-@st.cache_resource(show_spinner=False)
-def load_doodle_model():
-    candidates = [
-        "quickdraw_model.keras",
-        "src/quickdraw_model.keras",
-        "quickdraw_model.h5",
-        "src/quickdraw_model.h5",
-    ]
-    for p in candidates:
-        if Path(p).exists() and Path(p).is_file() and Path(p).stat().st_size > 0:
-            m = load_model(p)
-            return m, p
-    return None, None
-
-model, model_path = load_doodle_model()
-if model_path:
-    st.caption(f"Model loaded: `{model_path}`")
-else:
-    st.warning("Model file not found. Running in demo mode with mock predictions.")
-
-# -------------------
-# Title + helper text
-# -------------------
+# Title Section
 st.markdown(
     f"""
     <div style='text-align:center;'>
-        <h1 style='color:{PRIMARY_COLOR}; margin-bottom:0;'>🎨 Guess the Doodle – AI Game</h1>
-        <p style='font-size:18px; margin-top:6px;'>Draw something below, and let AI guess it in real-time!</p>
+        <h1 style='color:{PRIMARY_COLOR};'>🎨 Guess the Doodle – AI Game</h1>
+        <p style='font-size:18px;'>Draw something below, then hit <b>Predict</b> to see AI’s guesses!</p>
     </div>
     """,
     unsafe_allow_html=True
 )
 
-# -------------
-# UI: 2 columns
-# -------------
-left, right = st.columns([1,1.1])
-
-with left:
-    st.subheader("🖌️ Draw here")
-    canvas_result = st_canvas(
-        fill_color=BG_COLOR,
-        stroke_width=10,
-        stroke_color="black",
-        background_color="white",
-        width=280,
-        height=280,
-        drawing_mode="freedraw",
-        key="canvas",
-    )
-    col_a, col_b = st.columns(2)
-    with col_a:
-        if st.button("Clear", type="secondary"):
-            st.experimental_rerun()
-    with col_b:
-        st.caption("Tip: Try a 🐱 cat, 🚗 car, 🌳 tree, ☀️ sun")
-
-with right:
+# Drawing Canvas
+canvas_result = st_canvas(
+    fill_color=BG_COLOR, stroke_width=10, stroke_color="black",
+    background_color="white", width=280, height=280,
+    drawing_mode="freedraw", key="canvas",
+)
+
+# Predict Button
+do_pred = st.button("🔮 Predict")
+
+# Prediction Block
+if do_pred and canvas_result.image_data is not None:
+    img = Image.fromarray((canvas_result.image_data).astype("uint8")).convert("L")
+    img = img.resize((28, 28))
+    img_array = img_to_array(img) / 255.0
+    img_array = np.expand_dims(img_array, axis=0)
+
+    preds = model.predict(img_array)[0]
+    top_indices = preds.argsort()[-3:][::-1]
+
     st.markdown(f"<h2 style='color:{ACCENT_COLOR};'>🤖 AI’s Best Guesses:</h2>", unsafe_allow_html=True)
 
-    if canvas_result.image_data is not None:
-        # Preprocess drawing to 28x28 grayscale
-        img = Image.fromarray((canvas_result.image_data).astype("uint8")).convert("L")
-        img = img.resize((28, 28))
-        img_array = img_to_array(img) / 255.0
-        img_array = np.expand_dims(img_array, axis=0)
-
-        # Predict
-        if model is not None:
-            preds = model.predict(img_array, verbose=0)[0]
-        else:
-            # Mock predictions so the demo UI still works if no model
-            preds = np.random.dirichlet(np.ones(len(CLASSES)), size=1)[0]
-
-        top_indices = preds.argsort()[-3:][::-1]
-
-        # Emoji map
-        emoji_map = {
-            "cat":"🐱","dog":"🐶","car":"���","tree":"🌳",
-            "house":"🏠","sun":"☀️","airplane":"✈️",
-            "fish":"🐟","flower":"🌸","bird":"🐦"
-        }
-
-        # Confetti / balloons if high confidence on top-1
-        if preds[top_indices[0]] >= 0.80:
-            st.balloons()
-
-        # Show guesses
-        for idx in top_indices:
-            label = CLASSES[idx]
-            score = float(preds[idx])
-            emoji = emoji_map.get(label, "🎯")
-
-            st.markdown(
-                f"<h3 style='color:{PRIMARY_COLOR}; margin:6px 0 4px 0;'>{emoji} {label} – {score*100:.1f}%</h3>",
-                unsafe_allow_html=True
-            )
-            st.progress(score)
-
-# Footer hint
-st.success("💡 Try: cat 🐱, car 🚗, tree 🌳, sun ☀️ — short, simple strokes work best!")
+    for idx in top_indices:
+        emoji = emoji_map.get(CLASSES[idx], "🎯")
+        st.markdown(
+            f"<h3 style='color:{PRIMARY_COLOR};'>{emoji} {CLASSES[idx]} – {preds[idx]*100:.1f}%</h3>",
+            unsafe_allow_html=True
+        )
+        st.progress(float(preds[idx]))
+
+    st.balloons()
+
+# Class Legend
+st.markdown(
+    "<hr><h4>📖 Supported Classes:</h4>" +
+    " ".join([f"{emoji_map[c]} {c}" for c in CLASSES]),
+    unsafe_allow_html=True
+)
+
+# Tip
+st.success("💡 Tip: Try drawing a cat 🐱, car 🚗, tree 🌳, or sun ☀️!")