Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,121 +1,161 @@
+# src/streamlit_app.py
+
 import os
 import json
 import numpy as np
 import streamlit as st
 from PIL import Image
-
 import tensorflow as tf
-from tensorflow.keras.applications import ResNet50
-from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout
-from tensorflow.keras.models import Sequential
 
-# -----------------------------
+# ==========================================================
 # Page config
-# -----------------------------
-st.set_page_config(page_title="ResNet50 Image Predictor", page_icon="🧠", layout="centered")
+# ==========================================================
+st.set_page_config(
+    page_title="ResNet50 Image Predictor",
+    page_icon="🧠",
+    layout="centered",
+)
+
 st.title("🧠 ResNet50 Image Predictor")
-st.write("Loads ResNet50 architecture and then loads trained weights (HF-safe).")
+st.write("Classifies mushroom images using a trained ResNet50 model.")
 
-# -----------------------------
-# Paths (FIXED)
-# -----------------------------
+# ==========================================================
+# Paths (HF Space: repo root is /app, your files are in /app/src)
+# ==========================================================
+MODEL_PATH = "src/new_best3_resnet50.keras"
 WEIGHTS_PATH = "src/resnet50_weights_noBN3.h5"
 CLASS_NAMES_PATH = "src/class_names3.json"
 
-# -----------------------------
-# Debug info
-# -----------------------------
-with st.expander("🔍 Debug info (HuggingFace check)"):
-    try:
-        st.write("Files in src/:")
-        st.write(os.listdir("src"))
-        st.write("Weights exists:", os.path.exists(WEIGHTS_PATH))
-        st.write("Class names exists:", os.path.exists(CLASS_NAMES_PATH))
-        st.write("TF:", tf.__version__)
-    except Exception as e:
-        st.error(e)
+IMG_SIZE = (224, 224)
 
-# -----------------------------
-# Load class names
-# -----------------------------
-def load_class_names(path: str) -> list:
+# ==========================================================
+# Helpers
+# ==========================================================
+@st.cache_data
+def load_class_names(path: str):
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"Missing class names file: {path}")
     with open(path, "r") as f:
-        names = json.load(f)
-    if not isinstance(names, list) or len(names) == 0:
-        raise ValueError("class_names3.json must be a non-empty list of class names.")
-    return names
-
-# -----------------------------
-# Build architecture (MUST match training!)
-# -----------------------------
-def build_model(num_classes: int) -> tf.keras.Model:
-    base_model = ResNet50(weights="imagenet", include_top=False, input_shape=(224, 224, 3))
-    base_model.trainable = False
-
-    model = Sequential([
-        base_model,
-        GlobalAveragePooling2D(),
-        Dense(256, activation="relu"),
-        Dropout(0.5),
-        Dense(num_classes, activation="softmax"),
-    ])
-    return model
+        data = json.load(f)
+
+    # Accept either list ["a","b"] or dict {"0":"a","1":"b"}
+    if isinstance(data, dict):
+        # sort by numeric key if possible
+        try:
+            items = sorted(data.items(), key=lambda kv: int(kv[0]))
+            data = [v for _, v in items]
+        except Exception:
+            data = list(data.values())
+
+    if not isinstance(data, list) or len(data) == 0:
+        raise ValueError("class_names JSON must be a non-empty list (or dict).")
+
+    return data
+
 
-# -----------------------------
-# Load model + weights
-# -----------------------------
 @st.cache_resource
-def load_trained_model(weights_path: str, class_names_path: str):
+def load_model_and_weights(model_path: str, weights_path: str):
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Missing model file: {model_path}")
     if not os.path.exists(weights_path):
         raise FileNotFoundError(f"Missing weights file: {weights_path}")
-    if not os.path.exists(class_names_path):
-        raise FileNotFoundError(f"Missing class names file: {class_names_path}")
 
-    class_names = load_class_names(class_names_path)
-    model = build_model(num_classes=len(class_names))
+    # Load architecture/model (compile False is fine for inference)
+    model = tf.keras.models.load_model(model_path, compile=False)
 
-    # IMPORTANT: build the model by calling it once (creates weights shapes)
-    _ = model(np.zeros((1, 224, 224, 3), dtype=np.float32), training=False)
-
-    # Load weights from your file name EXACTLY
+    # Load trained weights (must match the loaded model)
     model.load_weights(weights_path)
 
-    return model, class_names
+    return model
+
+
+def preprocess_for_resnet50(pil_img: Image.Image) -> np.ndarray:
+    img = pil_img.convert("RGB").resize(IMG_SIZE)
+    x = np.array(img, dtype=np.float32)
+    x = np.expand_dims(x, axis=0)
+    x = tf.keras.applications.resnet50.preprocess_input(x)
+    return x
+
+
+# ==========================================================
+# Debug info (helpful on HuggingFace)
+# ==========================================================
+with st.expander("🔍 Debug info (HuggingFace check)"):
+    try:
+        st.write("Working directory:", os.getcwd())
+        st.write("Files in ./ :", os.listdir("."))
+        if os.path.exists("src"):
+            st.write("Files in src/:", os.listdir("src"))
+        else:
+            st.warning("Folder 'src' not found.")
+        st.write("MODEL_PATH exists:", os.path.exists(MODEL_PATH))
+        st.write("WEIGHTS_PATH exists:", os.path.exists(WEIGHTS_PATH))
+        st.write("CLASS_NAMES exists:", os.path.exists(CLASS_NAMES_PATH))
+        st.write("TensorFlow version:", tf.__version__)
+    except Exception as e:
+        st.exception(e)
 
+# ==========================================================
+# Load artifacts
+# ==========================================================
 try:
-    model, class_names = load_trained_model(WEIGHTS_PATH, CLASS_NAMES_PATH)
+    class_names = load_class_names(CLASS_NAMES_PATH)
+except Exception as e:
+    st.error("❌ Could not load class names.")
+    st.exception(e)
+    st.stop()
+
+# Show available classes
+with st.expander("🧪 Available mushroom classes (you can test these)"):
+    for i, name in enumerate(class_names):
+        st.write(f"**{i}** → {name}")
+
+try:
+    model = load_model_and_weights(MODEL_PATH, WEIGHTS_PATH)
     st.success("✅ Model + weights loaded successfully!")
 except Exception as e:
     st.error("❌ Model could not be loaded.")
     st.exception(e)
     st.stop()
 
-# -----------------------------
-# Image upload & prediction
-# -----------------------------
-uploaded_file = st.file_uploader("Upload a mushroom image", type=["jpg", "jpeg", "png", "webp"])
+# ==========================================================
+# Upload + predict
+# ==========================================================
+uploaded_file = st.file_uploader(
+    "Upload a mushroom image",
+    type=["jpg", "jpeg", "png", "webp"],
+)
 
 if uploaded_file is None:
     st.info("👆 Please upload an image to start prediction.")
-else:
-    img = Image.open(uploaded_file).convert("RGB")
-    st.image(img, caption="Uploaded image", use_container_width=True)
-
-    img_resized = img.resize((224, 224))
-    x = np.array(img_resized, dtype=np.float32)
-    x = np.expand_dims(x, axis=0)
-    x = tf.keras.applications.resnet50.preprocess_input(x)
-
-    preds = model.predict(x, verbose=0)[0]
-    idx = int(np.argmax(preds))
-    conf = float(np.max(preds))
-
-    st.subheader("✅ Prediction")
-    st.write(f"**Predicted class:** {class_names[idx]}")
-    st.write(f"**Confidence:** {conf:.4f}")
+    st.stop()
 
-    st.subheader("📊 Class probabilities")
-    # Show top 5
-    topk = np.argsort(preds)[::-1][:5]
-    for i in topk:
-        st.write(f"{class_names[i]}: {preds[i]:.4f}")
+# Show image
+img = Image.open(uploaded_file)
+st.image(img, caption="Uploaded image", use_container_width=True)
+
+# Predict
+x = preprocess_for_resnet50(img)
+preds = model.predict(x, verbose=0)[0]  # shape: (num_classes,)
+
+pred_idx = int(np.argmax(preds))
+pred_name = class_names[pred_idx] if pred_idx < len(class_names) else f"class_{pred_idx}"
+confidence = float(preds[pred_idx])
+
+st.subheader("✅ Prediction")
+st.write(f"**Predicted class index:** {pred_idx}")
+st.write(f"**Predicted mushroom:** 🍄 **{pred_name}**")
+st.write(f"**Confidence:** {confidence:.4f}")
+
+# Top-3
+st.subheader("📊 Top-3 predictions")
+top3_idx = np.argsort(preds)[::-1][:3]
+for i in top3_idx:
+    name = class_names[i] if i < len(class_names) else f"class_{i}"
+    st.write(f"**{name}** → {preds[i]:.4f}")
+
+# Full probs (optional)
+with st.expander("All class probabilities"):
+    for i, p in enumerate(preds):
+        name = class_names[i] if i < len(class_names) else f"class_{i}"
+        st.write(f"{i} → {name}: {p:.6f}")